diff --git a/ggml/src/ggml-hrx/CMakeLists.txt b/ggml/src/ggml-hrx/CMakeLists.txt index d4e79bc65456..b3eb70120f57 100644 --- a/ggml/src/ggml-hrx/CMakeLists.txt +++ b/ggml/src/ggml-hrx/CMakeLists.txt @@ -1,4 +1,5 @@ set(HRX_SOURCE_DIR "" CACHE PATH "Optional HRX source tree to build instead of using installed hrx and loomc packages") +find_package(Python3 REQUIRED COMPONENTS Interpreter) if(HRX_SOURCE_DIR) include(ExternalProject) @@ -9,12 +10,28 @@ if(HRX_SOURCE_DIR) message(FATAL_ERROR "HRX_SOURCE_DIR does not contain a CMakeLists.txt: ${HRX_SOURCE_DIR}") endif() + # Enforce the HRX compatibility floor before building dependencies; rerun configure when the pin or checker changes. + set(GGML_HRX_REVISION_CHECK "${CMAKE_CURRENT_SOURCE_DIR}/tools/check_hrx_revision.py") + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/hrx-minimum-commit.txt" "${GGML_HRX_REVISION_CHECK}") + execute_process( + COMMAND "${Python3_EXECUTABLE}" "${GGML_HRX_REVISION_CHECK}" "${HRX_SOURCE_DIR}" + RESULT_VARIABLE GGML_HRX_REVISION_RESULT + OUTPUT_VARIABLE GGML_HRX_REVISION_OUTPUT + ERROR_VARIABLE GGML_HRX_REVISION_ERROR + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT GGML_HRX_REVISION_RESULT EQUAL 0) + message(FATAL_ERROR "${GGML_HRX_REVISION_ERROR}") + endif() + message(STATUS "${GGML_HRX_REVISION_OUTPUT}") + set(GGML_HRX_PREFIX "${CMAKE_CURRENT_BINARY_DIR}/hrx") set(GGML_HRX_BUILD_DIR "${GGML_HRX_PREFIX}/src/ggml-hrx-deps-build") set(GGML_HRX_LIB "${GGML_HRX_BUILD_DIR}/libhrx/src/libhrx/${CMAKE_SHARED_LIBRARY_PREFIX}hrx${CMAKE_SHARED_LIBRARY_SUFFIX}") set(GGML_LOOMC_LIB "${GGML_HRX_BUILD_DIR}/loom/binding/c/${CMAKE_SHARED_LIBRARY_PREFIX}loomc${CMAKE_SHARED_LIBRARY_SUFFIX}") set(GGML_HRX_LOOM_LINK "${GGML_HRX_BUILD_DIR}/loom/src/loom/tools/loom-link/loom-link${CMAKE_EXECUTABLE_SUFFIX}") set(GGML_HRX_LOOM_FORMAT "${GGML_HRX_BUILD_DIR}/loom/src/loom/tools/loom-format/loom-format${CMAKE_EXECUTABLE_SUFFIX}") + set(GGML_HRX_IREE_BENCHMARK_LOOM "${GGML_HRX_BUILD_DIR}/loom/src/loom/tools/iree-benchmark-loom/iree-benchmark-loom${CMAKE_EXECUTABLE_SUFFIX}") set(GGML_HRX_DEPS_TARGET ggml-hrx-deps) set(GGML_HRX_CMAKE_ARGS @@ -41,9 +58,10 @@ if(HRX_SOURCE_DIR) SOURCE_DIR "${HRX_SOURCE_DIR}" PREFIX "${GGML_HRX_PREFIX}" CMAKE_ARGS ${GGML_HRX_CMAKE_ARGS} - BUILD_COMMAND ${CMAKE_COMMAND} --build . --target hrx loomc_shared loom_tools_loom-link_loom-link loom_tools_loom-format_loom-format --config ${CMAKE_BUILD_TYPE} + BUILD_ALWAYS TRUE + BUILD_COMMAND ${CMAKE_COMMAND} --build . --target hrx loomc_shared loom_tools_loom-link_loom-link loom_tools_loom-format_loom-format loom_tools_iree-benchmark-loom_iree-benchmark-loom --config ${CMAKE_BUILD_TYPE} INSTALL_COMMAND "" - BUILD_BYPRODUCTS "${GGML_HRX_LIB}" "${GGML_LOOMC_LIB}" "${GGML_HRX_LOOM_LINK}" "${GGML_HRX_LOOM_FORMAT}" + BUILD_BYPRODUCTS "${GGML_HRX_LIB}" "${GGML_LOOMC_LIB}" "${GGML_HRX_LOOM_LINK}" "${GGML_HRX_LOOM_FORMAT}" "${GGML_HRX_IREE_BENCHMARK_LOOM}" UPDATE_COMMAND "" ) @@ -64,10 +82,9 @@ else() find_package(loomc CONFIG REQUIRED) find_program(GGML_HRX_LOOM_LINK NAMES loom-link) find_program(GGML_HRX_LOOM_FORMAT NAMES loom-format) + find_program(GGML_HRX_IREE_BENCHMARK_LOOM NAMES iree-benchmark-loom) endif() -find_package(Python3 REQUIRED COMPONENTS Interpreter) - set(GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT "binary" CACHE STRING "Embedded Loom corpus source format: text or binary") set_property(CACHE GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT PROPERTY STRINGS text binary) if(GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT STREQUAL "binary") @@ -86,8 +103,30 @@ elseif(NOT GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT STREQUAL "text") message(FATAL_ERROR "Unsupported GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT: ${GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT}") endif() -set(GGML_HRX_QWEN_KERNEL_CORPUS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/kernel-corpus/kernels/qwen_moe") +set(GGML_HRX_QWEN_MOE_KERNEL_CORPUS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/kernel-corpus/kernels/qwen_moe") +set(GGML_HRX_QWEN_MOE_KERNEL_CORPUS_MANIFEST "${GGML_HRX_QWEN_MOE_KERNEL_CORPUS_DIR}/manifest.json") +set(GGML_HRX_HRX_KERNEL_CORPUS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/kernel-corpus/kernels/hrx") +set(GGML_HRX_HRX_KERNEL_CORPUS_MANIFEST "${GGML_HRX_HRX_KERNEL_CORPUS_DIR}/manifest.json") +set(GGML_HRX_QWEN_KERNEL_CORPUS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/kernel-corpus/kernels/qwen") set(GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST "${GGML_HRX_QWEN_KERNEL_CORPUS_DIR}/manifest.json") +set(GGML_HRX_LOOM_LIBS_KERNEL_CORPUS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/kernel-corpus/kernels/loom-libs") +set(GGML_HRX_LOOM_LIBS_KERNEL_CORPUS_MANIFEST "${GGML_HRX_LOOM_LIBS_KERNEL_CORPUS_DIR}/manifest.json") +set(GGML_HRX_KERNEL_CORPUS_MANIFESTS + "${GGML_HRX_QWEN_MOE_KERNEL_CORPUS_MANIFEST}" + "${GGML_HRX_HRX_KERNEL_CORPUS_MANIFEST}" + "${GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST}" + "${GGML_HRX_LOOM_LIBS_KERNEL_CORPUS_MANIFEST}" +) +set(GGML_HRX_KERNEL_CORPUS_MANIFEST_ARGS + --manifest "${GGML_HRX_QWEN_MOE_KERNEL_CORPUS_MANIFEST}" + --corpus-dir "${GGML_HRX_QWEN_MOE_KERNEL_CORPUS_DIR}" + --manifest "${GGML_HRX_HRX_KERNEL_CORPUS_MANIFEST}" + --corpus-dir "${GGML_HRX_HRX_KERNEL_CORPUS_DIR}" + --manifest "${GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST}" + --corpus-dir "${GGML_HRX_QWEN_KERNEL_CORPUS_DIR}" + --manifest "${GGML_HRX_LOOM_LIBS_KERNEL_CORPUS_MANIFEST}" + --corpus-dir "${GGML_HRX_LOOM_LIBS_KERNEL_CORPUS_DIR}" +) set(GGML_HRX_KERNEL_CORPUS_SOURCES_INC "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus-sources.inc") set(GGML_HRX_KERNEL_CORPUS_QWEN_INC "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus-qwen.inc") set(GGML_HRX_KERNEL_CORPUS_CATALOG_INC "${CMAKE_CURRENT_BINARY_DIR}/kernel-corpus-catalog.inc") @@ -103,14 +142,13 @@ add_custom_command( --source-output "${GGML_HRX_KERNEL_CORPUS_SOURCES_INC}" --corpus-output "${GGML_HRX_KERNEL_CORPUS_QWEN_INC}" --catalog-output "${GGML_HRX_KERNEL_CORPUS_CATALOG_INC}" - --manifest "${GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST}" - --corpus-dir "${GGML_HRX_QWEN_KERNEL_CORPUS_DIR}" + ${GGML_HRX_KERNEL_CORPUS_MANIFEST_ARGS} --source-format "${GGML_HRX_KERNEL_CORPUS_SOURCE_FORMAT}" ${GGML_HRX_KERNEL_CORPUS_TOOL_ARGS} --depfile "${GGML_HRX_KERNEL_CORPUS_DEPFILE}" DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/tools/generate_kernel_corpus.py" - "${GGML_HRX_QWEN_KERNEL_CORPUS_MANIFEST}" + ${GGML_HRX_KERNEL_CORPUS_MANIFESTS} ${GGML_HRX_KERNEL_CORPUS_TOOL_DEPENDS} DEPFILE "${GGML_HRX_KERNEL_CORPUS_DEPFILE}" VERBATIM @@ -148,6 +186,8 @@ ggml_add_backend_library(ggml-hrx dispatch/command-plan.h dispatch/command-program-bindings.cpp dispatch/command-program-bindings.h + dispatch/command-program-dump.cpp + dispatch/command-program-dump.h dispatch/command-program-diagnostics.cpp dispatch/command-program-diagnostics.h dispatch/command-program.cpp @@ -159,30 +199,69 @@ ggml_add_backend_library(ggml-hrx dispatch/dispatch.h dispatch/transient-allocator.cpp dispatch/transient-allocator.h - dispatch_registration/dispatch-add.cpp - dispatch_registration/dispatch-add.h - dispatch_registration/dispatch-gather-add.cpp - dispatch_registration/dispatch-gather-add.h - dispatch_registration/dispatch-llm-matmul.cpp - dispatch_registration/dispatch-llm-matmul.h - dispatch_registration/dispatch-llm-profiles.h - dispatch_registration/dispatch-llm-shapes.h - dispatch_registration/dispatch-qwen-attention-postprocess.cpp - dispatch_registration/dispatch-qwen-attention-postprocess.h - dispatch_registration/dispatch-qwen-flash-attention.cpp - dispatch_registration/dispatch-qwen-flash-attention.h - dispatch_registration/dispatch-qwen-matmul.cpp - dispatch_registration/dispatch-qwen-matmul.h - dispatch_registration/dispatch-moe-router.cpp - dispatch_registration/dispatch-moe-router.h - dispatch_registration/dispatch-qwen-preamble.cpp - dispatch_registration/dispatch-qwen-preamble.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 + dispatch_registration/common/dispatch-common.cpp + dispatch_registration/common/dispatch-common.h + dispatch_registration/common/dispatch-copy.cpp + dispatch_registration/common/dispatch-copy.h + dispatch_registration/common/dispatch-flash-attention.cpp + dispatch_registration/common/dispatch-flash-attention.h + dispatch_registration/common/dispatch-gated-mul-mat-id.cpp + dispatch_registration/common/dispatch-gated-mul-mat-id.h + dispatch_registration/common/dispatch-gated-mul-mat.cpp + dispatch_registration/common/dispatch-gated-mul-mat.h + dispatch_registration/common/dispatch-gather-add.cpp + dispatch_registration/common/dispatch-gather-add.h + dispatch_registration/common/dispatch-get-rows.cpp + dispatch_registration/common/dispatch-get-rows.h + dispatch_registration/common/dispatch-glu.cpp + dispatch_registration/common/dispatch-glu.h + dispatch_registration/common/dispatch-mul-mat-id.cpp + dispatch_registration/common/dispatch-mul-mat-id-common.h + dispatch_registration/common/dispatch-mul-mat-id.h + dispatch_registration/common/dispatch-mul-mat.cpp + dispatch_registration/common/dispatch-mul-mat-common.h + dispatch_registration/common/dispatch-mul-mat-weight-format.h + dispatch_registration/common/dispatch-mul-mat.h + dispatch_registration/common/dispatch-rope-set-rows.cpp + dispatch_registration/common/dispatch-rope-set-rows.h + dispatch_registration/common/dispatch-rope-utils.h + dispatch_registration/common/dispatch-rmsnorm.cpp + dispatch_registration/common/dispatch-rmsnorm.h + dispatch_registration/common/dispatch-scale.cpp + 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/dispatch-rmsnorm.cpp - dispatch_registration/dispatch-rmsnorm.h - dispatch_registration/dispatch-routed-ffn.cpp - dispatch_registration/dispatch-routed-ffn.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 + dispatch_registration/llm/dispatch-gated-delta-net.h + dispatch_registration/llm/dispatch-ssm-conv.cpp + dispatch_registration/llm/dispatch-ssm-conv.h + dispatch_registration/qwen/dispatch-llm-profiles.h + dispatch_registration/qwen/dispatch-llm-shapes.h + dispatch_registration/qwen/dispatch-moe-router.cpp + dispatch_registration/qwen/dispatch-moe-router.h + dispatch_registration/qwen/dispatch-qwen.cpp + dispatch_registration/qwen/dispatch-qwen.h + dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp + dispatch_registration/qwen/dispatch-qwen-attention-postprocess.h + dispatch_registration/qwen/dispatch-qwen-matmul.cpp + dispatch_registration/qwen/dispatch-qwen-matmul.h + dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp + dispatch_registration/qwen/dispatch-qwen-rmsnorm.h + dispatch_registration/qwen/dispatch-routed-ffn.cpp + dispatch_registration/qwen/dispatch-routed-ffn.h status.h graph/graph.cpp graph/graph-diagnostics.cpp diff --git a/ggml/src/ggml-hrx/README.md b/ggml/src/ggml-hrx/README.md new file mode 100644 index 000000000000..cf00e3131c6f --- /dev/null +++ b/ggml/src/ggml-hrx/README.md @@ -0,0 +1,17 @@ +# HRX Backend + +## `hrx-system` backend compatibility + +`hrx-minimum-commit.txt` records the minimum required commit in ROCm/hrx-system. It contains one full Git commit SHA. The backend's C++ APIs, Loom syntax, and numerical behavior can all depend on that revision. When `GGML_HRX=ON` and `HRX_SOURCE_DIR` is set, CMake requires the HRX checkout's HEAD to equal or descend from the recorded commit. + +For builds using installed packages, check the source revision used to produce both HRX and Loom against this floor before validating a repair. From the staging checkout: + +```sh +python3 llama.cpp/ggml/src/ggml-hrx/tools/check_hrx_revision.py hrx-system +``` + +### Updating the minimum + +Raise the minimum in the same llama.cpp commit that introduces a newer HRX dependency. This includes new API calls, new Loom syntax, and removing a numerical workaround after an HRX compiler or runtime fix. Unrelated repairs do not require a bump. + +Choose the earliest demonstrated required commit that retains the previous floor in its ancestry. Do not replace the floor with the latest tested revision merely because it is available. Record the dependency and its HRX commit or PR in the change's validation evidence. A minimum is not a promise that every future HRX revision is compatible. diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h index 0d5472f01cd8..3d510b310b59 100644 --- a/ggml/src/ggml-hrx/backend-context.h +++ b/ggml/src/ggml-hrx/backend-context.h @@ -63,6 +63,7 @@ struct ggml_backend_hrx_context { ggml::hrx::TransientArena transient_arena; ggml::hrx::HostTransferManager host_transfers; ggml::hrx::HostWeightCache host_weights; + ggml::hrx::GraphReplayStreamState graph_replay_state; std::string name; }; diff --git a/ggml/src/ggml-hrx/benchmarks/README.md b/ggml/src/ggml-hrx/benchmarks/README.md new file mode 100644 index 000000000000..58b8e83db874 --- /dev/null +++ b/ggml/src/ggml-hrx/benchmarks/README.md @@ -0,0 +1,199 @@ +# HRX Loom Benchmarks + +This directory contains model-scoped Loom benchmarks for HRX kernels. Benchmark sources stay separate from the production kernel corpus: they declare the kernels they need with `kernel.decl`, and the runner links those declarations against the production `.loom` files before benchmarking. + +Use these benchmarks to measure production kernels with `iree-benchmark-loom` while keeping model-shaped benchmark cases out of the embedded kernel catalog. Each model has one `.loom` file under `loom/`, and each workload shape has a sidecar manifest next to it, for example `llama32_3b_f16.pp512.json`. + +Generated model benchmarks are deduplicated by kernel shape. The sidecar manifest keeps a compact `dispatches` list with the generated benchmark name, kernel, compile/runtime parameters, source files, and `count`, so benchmark results can be weighted back to the full model without materializing one check case for every duplicate dispatch. + +Generated model benchmarks are performance fixtures, not numerical validators. They launch the model-shaped kernels with synthetic buffers and intentionally omit `check.expect.*` assertions; backend model smoke tests and the kernel corpus checks remain responsible for numerical validation. + +## Generating Llama Benchmarks + +First dump the HRX command programs for each model run: + +```sh +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-llama32-pp256-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/llama-3.2/Llama-3.2-3B-Instruct-F16.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 64 \ + --ubatch-size 64 \ + --repetitions 1 \ + --no-warmup \ + --output jsonl \ + --n-prompt 256 \ + --n-gen 0 \ + --n-depth 0 + +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-llama32-pp512-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/llama-3.2/Llama-3.2-3B-Instruct-F16.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 64 \ + --ubatch-size 64 \ + --repetitions 1 \ + --no-warmup \ + --output jsonl \ + --n-prompt 512 \ + --n-gen 0 \ + --n-depth 0 + +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-llama32-tg8-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/llama-3.2/Llama-3.2-3B-Instruct-F16.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 64 \ + --ubatch-size 64 \ + --repetitions 1 \ + --no-warmup \ + --output jsonl \ + --n-prompt 0 \ + --n-gen 8 \ + --n-depth 0 +``` + +For the Llama 3.2 1B Q4_K_XL decode benchmark, dump the tg32 command program with: + +```sh +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-llama32-1b-tg32-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/lemonade/llamacpp-gguf-models/unsloth_Llama-3.2-1B-Instruct-GGUF/Llama-3.2-1B-Instruct-UD-Q4_K_XL.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 512 \ + --ubatch-size 512 \ + --repetitions 1 \ + --output jsonl \ + --n-prompt 0 \ + --n-gen 32 \ + --n-depth 0 +``` + +Then generate the shared model benchmark file and per-scenario sidecars: + +```sh +ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py \ + --model llama32_3b_f16 \ + --scenario-dump pp256=/tmp/hrx-llama32-pp256-dumps \ + --scenario-dump pp512=/tmp/hrx-llama32-pp512-dumps \ + --scenario-dump tg8=/tmp/hrx-llama32-tg8-dumps + +ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py \ + --model llama32_1b_q4_k_xl \ + --scenario tg32 \ + --dump-dir /tmp/hrx-llama32-1b-tg32-dumps +``` + +Review the generated invoked-kernel set with: + +```sh +jq '.kernel_counts' ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16..json + +ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.sh \ + --model llama32_3b_f16 \ + --scenario \ + --build-dir \ + --list +``` + +Run the generated benchmarks with: + +```sh +ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.sh \ + --model llama32_3b_f16 \ + --scenario \ + --build-dir \ + --output-dir /home/rsuderman/codex/project-workspaces/llama.cpp/gates/hrx-loom-benchmarks/llama32- +``` + +Use `--dry-run-only` to stop after planning, `--list` to inspect the generated benchmark set, or `--benchmark @name` to run a single invocation. Set `LOOM_LINK` or `IREE_BENCHMARK_LOOM` to override tool discovery. Set `DEVICE` to override the default `amdgpu` HAL device. + +The runner applies the same workload-argument specialization used by the HRX runtime after `loom-link`, and benchmarks the resulting `linked.runtime-specialized.loom` source by default. Use `--no-runtime-specialization` to benchmark the raw linked source. + +`loom-link` also receives each captured compile parameter as `--config==`. Some model kernels select templates from compile-time config values, so passing those configs during link keeps standalone benchmarks aligned with the HRX runtime path. + +Use `--continue-on-failure` to record failed or timed-out standalone cases and keep running the rest of the model benchmark set. Use `--benchmark-timeout-sec` to bound each link, dry-run, and benchmark command. Each benchmark directory keeps `loom-link`, dry-run, and benchmark command lines plus stdout/stderr so failures can be reproduced directly. + +Use `--profile-final-batch=false`, `--input-ring-count`, or repeated `--benchmark-extra-arg` flags when isolating `iree-benchmark-loom` behavior from production runtime behavior. After a run, summarize the weighted model time with: + +```sh +ggml/src/ggml-hrx/tools/benchmarks/summarize-model-benchmarks.py \ + /home/rsuderman/codex/project-workspaces/llama.cpp/gates/hrx-loom-benchmarks/llama32-pp512/results.jsonl +``` + +The summary uses `operation_timing_ns.p50 * count` by default and writes `summary.json` plus `summary.md` next to the runner results. + +To map likely fusion opportunities from a command-program dump and the weighted benchmark summary: + +```sh +ggml/src/ggml-hrx/tools/benchmarks/analyze-model-fusion-adjacency.py \ + --dump-dir /tmp/hrx-llama32-tg8-dumps \ + --scenario-manifest ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json \ + --summary-json /home/rsuderman/codex/project-workspaces/llama.cpp/gates/hrx-loom-benchmarks/llama32-tg8/summary.json \ + --output-json /home/rsuderman/codex/project-workspaces/llama.cpp/gates/hrx-loom-benchmarks/llama32-tg8/fusion-adjacency.json \ + --output-md /home/rsuderman/codex/project-workspaces/llama.cpp/gates/hrx-loom-benchmarks/llama32-tg8/fusion-adjacency.md +``` + +The adjacency report uses transient value producer-consumer edges for true data dependencies and sequential cache-update windows for side-effect patterns such as RoPE, SET_ROWS, and decode flash attention. + +## Generating Qwen 30B Benchmarks + +Dump the HRX command programs with the Qwen 30B shard and the same pp256, pp512, and tg8 shapes used for model benchmarking: + +```sh +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-qwen30b-pp256-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/qwen-30b/qwen3-30b-a3b-q4_k_m-00001-of-00020.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 64 \ + --ubatch-size 64 \ + --repetitions 1 \ + --no-warmup \ + --output jsonl \ + --n-prompt 256 \ + --n-gen 0 \ + --n-depth 0 + +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-qwen30b-pp512-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/qwen-30b/qwen3-30b-a3b-q4_k_m-00001-of-00020.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 64 \ + --ubatch-size 64 \ + --repetitions 1 \ + --no-warmup \ + --output jsonl \ + --n-prompt 512 \ + --n-gen 0 \ + --n-depth 0 + +GGML_HRX_DUMP_COMMAND_PROGRAM_DIR=/tmp/hrx-qwen30b-tg8-dumps \ +/bin/llama-bench \ + --model /home/rsuderman/Downloads/gguf/qwen-30b/qwen3-30b-a3b-q4_k_m-00001-of-00020.gguf \ + --device HRX0 \ + --n-gpu-layers -1 \ + --batch-size 64 \ + --ubatch-size 64 \ + --repetitions 1 \ + --no-warmup \ + --output jsonl \ + --n-prompt 0 \ + --n-gen 8 \ + --n-depth 0 +``` + +Then regenerate the shared Qwen benchmark source and sidecars: + +```sh +ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py \ + --model qwen3_30b_a3b_q4_k_m \ + --scenario-dump pp256=/tmp/hrx-qwen30b-pp256-dumps \ + --scenario-dump pp512=/tmp/hrx-qwen30b-pp512-dumps \ + --scenario-dump tg8=/tmp/hrx-qwen30b-tg8-dumps +``` diff --git a/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json new file mode 100644 index 000000000000..e71bf8fb39cf --- /dev/null +++ b/ggml/src/ggml-hrx/benchmarks/loom/llama32_3b_f16.tg8.json @@ -0,0 +1,626 @@ +{ + "command_count": 451, + "dispatch_count": 14, + "dispatches": [ + { + "benchmark": "@llama32_3b_f16_tg8_000_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "3072", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 3072, + "row_count": 128256, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_001_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "3072", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999975e-06" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 57, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_002_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "3072", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 56, + "integer_parameters": { + "input_size": 3072, + "output_size": 3072, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_003_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "1024", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 56, + "integer_parameters": { + "input_size": 3072, + "output_size": 1024, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_004_ggml_rope_f32", + "compile_parameters": { + "ggml.rope_f32.head_count": "24", + "ggml.rope_f32.head_size": "128", + "ggml.rope_f32.mode": "0", + "ggml.rope_f32.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 28, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rope_f32", + "library_sources": [ + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/rope/rope_f32.loom" + ], + "sources": [ + "ops/rope/rope_f32.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rope_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_005_ggml_rope_set_rows_f32", + "compile_parameters": { + "ggml.rope_set_rows_f32.head_count": "8", + "ggml.rope_set_rows_f32.head_size": "128", + "ggml.rope_set_rows_f32.mode": "0", + "ggml.rope_set_rows_f32.output_format": "16", + "ggml.rope_set_rows_f32.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 28, + "integer_parameters": { + "cache_row_count": 256, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rope_set_rows_f32", + "library_sources": [ + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rope/rope_set_rows_f32.loom" + ], + "sources": [ + "ops/rope/rope_set_rows_f32.loom", + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_rope_set_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_006_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 28, + "integer_parameters": { + "cache_row_count": 256, + "hidden_size": 1024, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_007_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "compile_parameters": { + "ggml.flash_attention.decode.key_value_token_capacity": "256", + "ggml.flash_attention.key_value_head_count": "8", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 28, + "integer_parameters": { + "key_value_token_count": 256 + }, + "kernel": "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "library_sources": [ + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ], + "primary_sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom" + ], + "sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ], + "symbol": "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_008_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 55, + "integer_parameters": { + "element_count": 3072 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_009_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "8192", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 56, + "integer_parameters": { + "input_size": 3072, + "output_size": 8192, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_010_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 28, + "integer_parameters": { + "element_count": 8192 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_011_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "3072", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 28, + "integer_parameters": { + "input_size": 8192, + "output_size": 3072, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_012_ggml_gather_add_f32", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/hrx", + "count": 1, + "integer_parameters": { + "hidden_size": 3072, + "output_token_count": 1, + "source_token_count": 1 + }, + "kernel": "hrx:ggml_gather_add_f32", + "library_sources": [], + "primary_sources": [ + "gather_add_f32.loom" + ], + "sources": [ + "gather_add_f32.loom" + ], + "symbol": "ggml_gather_add_f32", + "workload_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@llama32_3b_f16_tg8_013_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "128256", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "input_size": 3072, + "output_size": 128256, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + } + ], + "generated_count": 14, + "kernel_counts": { + "hrx:ggml_gather_add_f32": 1, + "loom_libs:ggml_binary_f32": 83, + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8": 28, + "loom_libs:ggml_get_rows_f32": 1, + "loom_libs:ggml_mul_mat_f32_f32_decode_wave64": 197, + "loom_libs:ggml_rmsnorm_binary_f32": 57, + "loom_libs:ggml_rope_f32": 28, + "loom_libs:ggml_rope_set_rows_f32": 28, + "loom_libs:ggml_set_rows": 28 + }, + "loom_source": "benchmarks/loom/llama32_3b_f16.loom", + "model": "llama32_3b_f16", + "scenario": "tg8", + "schema": "ggml-hrx-model-loom-benchmarks-v2" +} diff --git a/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json new file mode 100644 index 000000000000..36d40a7130f8 --- /dev/null +++ b/ggml/src/ggml-hrx/benchmarks/loom/qwen3_30b_a3b_q4_k_m.tg8.json @@ -0,0 +1,1051 @@ +{ + "command_count": 304, + "dispatch_count": 21, + "dispatches": [ + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_000_qwen_attention_context_base_capture", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/qwen", + "count": 1, + "integer_parameters": {}, + "kernel": "qwen:qwen_attention_context_base_capture", + "library_sources": [], + "primary_sources": [ + "attention_state_initialize.loom" + ], + "sources": [ + "attention_state_initialize.loom" + ], + "symbol": "qwen_attention_context_base_capture", + "workload_parameters": [] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_001_qwen_attention_metadata", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 1, + "integer_parameters": { + "context_capacity": 256, + "token_count": 1 + }, + "kernel": "qwen3_moe:qwen_attention_metadata", + "library_sources": [], + "primary_sources": [ + "../qwen/attention_metadata.loom" + ], + "sources": [ + "../qwen/attention_metadata.loom" + ], + "symbol": "qwen_attention_metadata", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_002_ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "2048", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "4", + "ggml.get_rows_rmsnorm.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 2048, + "row_count": 151936, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "library_sources": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/normalization/publish_q8.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom" + ], + "sources": [ + "ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom", + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/normalization/publish_q8.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_003_qwen3_moe_attention_qkv_postprocess_fused_decode", + "compile_parameters": { + "qwen3_moe.attention.head_size": "128", + "qwen3_moe.attention.key_value_size": "512", + "qwen3_moe.attention.query_size": "4096", + "qwen3_moe.attention.rope_mscale": "1", + "qwen3_moe.attention.value_uses_q6": "1", + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.model.rms_epsilon": "0.000001", + "qwen3_moe.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 24, + "integer_parameters": { + "cache_row_count": 256, + "token_count": 1 + }, + "kernel": "qwen3_moe:qwen3_moe_attention_qkv_postprocess_fused_decode", + "library_sources": [ + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "primary_sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom" + ], + "sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_004_ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "compile_parameters": { + "ggml.flash_attention.attention_scale": "0.0883883461", + "ggml.flash_attention.decode.key_value_token_capacity": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.qk_head_size": "128", + "ggml.flash_attention.query_head_count": "32", + "ggml.flash_attention.value_head_size": "128" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "key_value_token_count": 256 + }, + "kernel": "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "library_sources": [ + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ], + "primary_sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom" + ], + "sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ], + "symbol": "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_005_qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "compile_parameters": { + "qwen3_moe.dense_quantized.input_size": "4096", + "qwen3_moe.dense_quantized.output_accumulation": "1", + "qwen3_moe.dense_quantized.output_size": "2048", + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.model.rms_epsilon": "0.000001", + "qwen3_moe.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 48, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "library_sources": [ + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "primary_sources": [ + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom" + ], + "sources": [ + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_006_qwen3_moe_router_projection_top8_fused_decode_f32", + "compile_parameters": { + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.router.expert_count": "128", + "qwen3_moe.router.route_count": "8", + "qwen3_moe.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 48, + "integer_parameters": { + "route_id_stride": 128, + "token_count": 1 + }, + "kernel": "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", + "library_sources": [ + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" + ], + "primary_sources": [ + "qwen3_moe/router_projection_top8_fused_f32.loom" + ], + "sources": [ + "qwen3_moe/router_projection_top8_fused_f32.loom", + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" + ], + "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_007_ggml_mul_mat_id_vector_pair_binary_publish_f32", + "compile_parameters": { + "ggml.mul_mat_id.vector.binary_op": "4", + "ggml.mul_mat_id.vector.expert_count": "128", + "ggml.mul_mat_id.vector.input_format": "9", + "ggml.mul_mat_id.vector.input_size": "2048", + "ggml.mul_mat_id.vector.lhs_weight_format": "4", + "ggml.mul_mat_id.vector.output_size": "768", + "ggml.mul_mat_id.vector.publish_format": "0", + "ggml.mul_mat_id.vector.rhs_weight_format": "4", + "ggml.mul_mat_id.vector.route_count": "8", + "ggml.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 24, + "integer_parameters": { + "expert_count": 128, + "output_size": 768, + "route_count": 8, + "route_stride": 128, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_f32", + "library_sources": [ + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/matmul_id/vector_pair_binary.loom" + ], + "sources": [ + "ops/matmul_id/vector_pair_binary.loom", + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_mul_mat_id_vector_pair_binary_publish_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_008_ggml_mul_mat_id_vector_weighted_wave64_publish_q8", + "compile_parameters": { + "ggml.mul_mat_id.vector.expert_count": "128", + "ggml.mul_mat_id.vector.input_format": "0", + "ggml.mul_mat_id.vector.input_size": "768", + "ggml.mul_mat_id.vector.output_size": "2048", + "ggml.mul_mat_id.vector.publish_format": "9", + "ggml.mul_mat_id.vector.route_count": "8", + "ggml.mul_mat_id.vector.weight_format": "6", + "ggml.workload.token_capacity": "1", + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.model.rms_epsilon": "0.000001" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 23, + "integer_parameters": { + "expert_count": 128, + "input_size": 768, + "output_size": 2048, + "route_count": 8, + "route_id_stride": 128, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_id_vector_weighted_wave64_publish_q8", + "library_sources": [ + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ], + "primary_sources": [ + "ops/matmul_id/vector_weighted.loom" + ], + "sources": [ + "ops/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ], + "symbol": "ggml_mul_mat_id_vector_weighted_wave64_publish_q8", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_009_qwen3_moe_attention_qkv_postprocess_fused_decode", + "compile_parameters": { + "qwen3_moe.attention.head_size": "128", + "qwen3_moe.attention.key_value_size": "512", + "qwen3_moe.attention.query_size": "4096", + "qwen3_moe.attention.rope_mscale": "1", + "qwen3_moe.attention.value_uses_q6": "0", + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.model.rms_epsilon": "0.000001", + "qwen3_moe.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 24, + "integer_parameters": { + "cache_row_count": 256, + "token_count": 1 + }, + "kernel": "qwen3_moe:qwen3_moe_attention_qkv_postprocess_fused_decode", + "library_sources": [ + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "primary_sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom" + ], + "sources": [ + "qwen3_moe/attention_qkv_postprocess_fused.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_010_ggml_mul_mat_id_vector_pair_binary_publish_q8", + "compile_parameters": { + "ggml.mul_mat_id.vector.binary_op": "4", + "ggml.mul_mat_id.vector.expert_count": "128", + "ggml.mul_mat_id.vector.input_format": "9", + "ggml.mul_mat_id.vector.input_size": "2048", + "ggml.mul_mat_id.vector.lhs_weight_format": "4", + "ggml.mul_mat_id.vector.output_size": "768", + "ggml.mul_mat_id.vector.publish_format": "9", + "ggml.mul_mat_id.vector.rhs_weight_format": "4", + "ggml.mul_mat_id.vector.route_count": "8", + "ggml.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 24, + "integer_parameters": { + "expert_count": 128, + "output_size": 768, + "route_count": 8, + "route_stride": 128, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_q8", + "library_sources": [ + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/matmul_id/vector_pair_binary.loom" + ], + "sources": [ + "ops/matmul_id/vector_pair_binary.loom", + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_mul_mat_id_vector_pair_binary_publish_q8", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_011_ggml_mul_mat_id_vector_weighted_wave32_publish_q8", + "compile_parameters": { + "ggml.mul_mat_id.vector.expert_count": "128", + "ggml.mul_mat_id.vector.input_format": "9", + "ggml.mul_mat_id.vector.input_size": "768", + "ggml.mul_mat_id.vector.output_size": "2048", + "ggml.mul_mat_id.vector.publish_format": "9", + "ggml.mul_mat_id.vector.route_count": "8", + "ggml.mul_mat_id.vector.weight_format": "4", + "ggml.workload.token_capacity": "1", + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.model.rms_epsilon": "0.000001" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 24, + "integer_parameters": { + "expert_count": 128, + "input_size": 768, + "output_size": 2048, + "route_count": 8, + "route_id_stride": 128, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_id_vector_weighted_wave32_publish_q8", + "library_sources": [ + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ], + "primary_sources": [ + "ops/matmul_id/vector_weighted.loom" + ], + "sources": [ + "ops/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ], + "symbol": "ggml_mul_mat_id_vector_weighted_wave32_publish_q8", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_012_ggml_moe_build_expert_table", + "compile_parameters": { + "ggml.moe_routing.descriptor_expert_mask": "511", + "ggml.moe_routing.descriptor_partition_shift": "9", + "ggml.moe_routing.descriptor_row_count_shift": "15", + "ggml.moe_routing.expert_count": "128", + "ggml.moe_routing.partition_workgroup_size": "512", + "ggml.moe_routing.route_count": "8", + "ggml.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "expert_count": 128, + "route_count": 8, + "route_stride": 8, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_moe_build_expert_table", + "library_sources": [], + "primary_sources": [ + "ops/moe_routing_tables.loom" + ], + "sources": [ + "ops/moe_routing_tables.loom" + ], + "symbol": "ggml_moe_build_expert_table", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_013_ggml_moe_build_expert_partition_table", + "compile_parameters": { + "ggml.moe_routing.descriptor_expert_mask": "511", + "ggml.moe_routing.descriptor_partition_shift": "9", + "ggml.moe_routing.descriptor_row_count_shift": "15", + "ggml.moe_routing.expert_count": "128", + "ggml.moe_routing.partition_workgroup_size": "512", + "ggml.moe_routing.route_count": "8", + "ggml.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "expert_count": 128, + "route_count": 8, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_moe_build_expert_partition_table", + "library_sources": [], + "primary_sources": [ + "ops/moe_routing_tables.loom" + ], + "sources": [ + "ops/moe_routing_tables.loom" + ], + "symbol": "ggml_moe_build_expert_partition_table", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_014_ggml_mul_mat_id_skinny_input_f32_publish_f32", + "compile_parameters": { + "ggml.mul_mat_id.expert_count": "128", + "ggml.mul_mat_id.input_route_count": "8", + "ggml.mul_mat_id.input_size": "768", + "ggml.mul_mat_id.output_size": "2048", + "ggml.mul_mat_id.route_count": "8", + "ggml.mul_mat_id.weight_format": "6", + "ggml.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_id_skinny_input_f32_publish_f32", + "library_sources": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/matmul_id/skinny_input_f32_publish_f32.loom" + ], + "sources": [ + "ops/matmul_id/skinny_input_f32_publish_f32.loom", + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_id_skinny_input_f32_publish_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_015_ggml_binary_bc_f32", + "compile_parameters": { + "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.src0_broadcast_dim2": "0", + "ggml.binary_bc_f32.src0_broadcast_dim3": "0", + "ggml.binary_bc_f32.src1_broadcast_dim0": "1", + "ggml.binary_bc_f32.src1_broadcast_dim1": "0", + "ggml.binary_bc_f32.src1_broadcast_dim2": "0", + "ggml.binary_bc_f32.src1_broadcast_dim3": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 16384, + "ne0": 2048, + "ne1": 8, + "ne2": 1, + "ne3": 1, + "src0_element_count": 16384, + "src1_element_count": 8 + }, + "kernel": "loom_libs:ggml_binary_bc_f32", + "library_sources": [ + "motifs/elementwise/binary_apply.loom", + "motifs/elementwise/broadcast_index.loom" + ], + "primary_sources": [ + "ops/elementwise/binary_bc_f32.loom" + ], + "sources": [ + "ops/elementwise/binary_bc_f32.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/elementwise/broadcast_index.loom" + ], + "symbol": "ggml_binary_bc_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + }, + { + "name": "ne0", + "type": "index" + }, + { + "name": "ne1", + "type": "index" + }, + { + "name": "ne2", + "type": "index" + }, + { + "name": "ne3", + "type": "index" + }, + { + "name": "src0_element_count", + "type": "index" + }, + { + "name": "src1_element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_016_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.ne0": "2048", + "ggml.binary_f32.ne1": "1", + "ggml.binary_f32.ne2": "1", + "ggml.binary_f32.op": "0", + "ggml.binary_f32.src0_span": "2048", + "ggml.binary_f32.src0_stride1": "16384", + "ggml.binary_f32.src0_stride2": "16384", + "ggml.binary_f32.src0_stride3": "16384", + "ggml.binary_f32.src1_span": "2048", + "ggml.binary_f32.src1_stride1": "16384", + "ggml.binary_f32.src1_stride2": "16384", + "ggml.binary_f32.src1_stride3": "16384" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 2048 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/elementwise/binary_f32.loom" + ], + "sources": [ + "ops/elementwise/binary_f32.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_017_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.ne0": "2048", + "ggml.binary_f32.ne1": "1", + "ggml.binary_f32.ne2": "1", + "ggml.binary_f32.op": "0", + "ggml.binary_f32.src0_span": "2048", + "ggml.binary_f32.src0_stride1": "2048", + "ggml.binary_f32.src0_stride2": "2048", + "ggml.binary_f32.src0_stride3": "2048", + "ggml.binary_f32.src1_span": "2048", + "ggml.binary_f32.src1_stride1": "16384", + "ggml.binary_f32.src1_stride2": "16384", + "ggml.binary_f32.src1_stride3": "16384" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 6, + "integer_parameters": { + "element_count": 2048 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/elementwise/binary_f32.loom" + ], + "sources": [ + "ops/elementwise/binary_f32.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_018_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.ne0": "2048", + "ggml.binary_f32.ne1": "1", + "ggml.binary_f32.ne2": "1", + "ggml.binary_f32.op": "0", + "ggml.binary_f32.src0_span": "2048", + "ggml.binary_f32.src0_stride1": "2048", + "ggml.binary_f32.src0_stride2": "2048", + "ggml.binary_f32.src0_stride3": "2048", + "ggml.binary_f32.src1_span": "2048", + "ggml.binary_f32.src1_stride1": "2048", + "ggml.binary_f32.src1_stride2": "2048", + "ggml.binary_f32.src1_stride3": "2048" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 2048 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/elementwise/binary_f32.loom" + ], + "sources": [ + "ops/elementwise/binary_f32.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_019_qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "compile_parameters": { + "ggml.quantize_q8_1_x4.group_capacity": "16", + "qwen3_moe.model.hidden_size": "2048", + "qwen3_moe.model.rms_epsilon": "0.000001", + "qwen3_moe.workload.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 1, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "library_sources": [ + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "primary_sources": [ + "qwen3_moe/attention_prepare_quantized.loom" + ], + "sources": [ + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "qwen3_moe/model_config.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen3_30b_a3b_q4_k_m_tg8_020_ggml_linear_q6k_q8_1_x4", + "compile_parameters": { + "ggml.linear_q6k_q8_1_x4.output_capacity": "151936", + "ggml.linear_q6k_q8_1_x4.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/qwen_moe", + "count": 1, + "integer_parameters": { + "input_size": 2048, + "output_size": 151936, + "token_count": 1 + }, + "kernel": "qwen3_moe:ggml_linear_q6k_q8_1_x4", + "library_sources": [ + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "primary_sources": [ + "ggml/linear_q6k_q8_1_x4.loom" + ], + "sources": [ + "ggml/linear_q6k_q8_1_x4.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" + ], + "symbol": "ggml_linear_q6k_q8_1_x4", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + } + ], + "generated_count": 21, + "kernel_counts": { + "loom_libs:ggml_binary_bc_f32": 1, + "loom_libs:ggml_binary_f32": 8, + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8": 48, + "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16": 1, + "loom_libs:ggml_moe_build_expert_partition_table": 1, + "loom_libs:ggml_moe_build_expert_table": 1, + "loom_libs:ggml_mul_mat_id_skinny_input_f32_publish_f32": 1, + "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_f32": 24, + "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_q8": 24, + "loom_libs:ggml_mul_mat_id_vector_weighted_wave32_publish_q8": 24, + "loom_libs:ggml_mul_mat_id_vector_weighted_wave64_publish_q8": 23, + "qwen3_moe:ggml_linear_q6k_q8_1_x4": 1, + "qwen3_moe:qwen3_moe_attention_qkv_postprocess_fused_decode": 48, + "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8": 48, + "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4": 1, + "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32": 48, + "qwen3_moe:qwen_attention_metadata": 1, + "qwen:qwen_attention_context_base_capture": 1 + }, + "loom_source": "benchmarks/loom/qwen3_30b_a3b_q4_k_m.loom", + "model": "qwen3_30b_a3b_q4_k_m", + "scenario": "tg8", + "schema": "ggml-hrx-model-loom-benchmarks-v2" +} diff --git a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp index eac956c644ac..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) { @@ -76,11 +85,13 @@ bool CommandPlanMetadata::append_generated_resource(CommandPlanGeneratedResource bool CommandPlanMetadata::append_alternate_value(CommandPlanAlternateValue alternate, Status & status) { for (const CommandPlanAlternateValue & existing : alternate_values_) { - if (existing.graph_value == alternate.graph_value) { + if (existing.graph_value == alternate.graph_value && existing.type == alternate.type && + existing.byte_count == alternate.byte_count) { if (alternate_value_matches(existing, alternate)) { return true; } - status.log("conflicting alternate value for graph value %d", alternate.graph_value.value); + status.log("conflicting alternate value for graph value %d type %d byte_count %zu", + alternate.graph_value.value, static_cast(alternate.type), alternate.byte_count); return false; } } @@ -124,11 +135,12 @@ const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(Valu const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(ValueId graph_value, ggml_type type, size_t byte_count) const { - const CommandPlanAlternateValue * alternate = find_alternate_value(graph_value); - if (alternate == nullptr || alternate->type != type || alternate->byte_count != byte_count) { - return nullptr; + for (const CommandPlanAlternateValue & alternate : alternate_values_) { + if (alternate.graph_value == graph_value && alternate.type == type && alternate.byte_count == byte_count) { + return &alternate; + } } - return alternate; + return nullptr; } const CommandPlanMoeRoutingBundle * CommandPlanMetadata::find_moe_routing_bundle(ValueId route_ids) const { diff --git a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h index f6d469c2cdaa..eeb8238690ad 100644 --- a/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h +++ b/ggml/src/ggml-hrx/dispatch/command-plan-metadata.h @@ -16,6 +16,8 @@ namespace ggml::hrx { enum class GeneratedResourceRole { MoeExpertTable, MoePartitionTable, + F16K16Major, + Conv4Edges, }; enum class CommandPlanResourceMetadataKind { @@ -84,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; @@ -107,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, @@ -124,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-bindings.cpp b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp index 1a3b1a76737e..741a41c9c17f 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.cpp @@ -28,7 +28,8 @@ CommandProgramBindings CommandProgramBindings::from_value_map(const ValueMap & v continue; } bindings.push_back({ value->id, buffer->buffer, buffer->offset, buffer->length, buffer->identity, - buffer->generation, buffer->capacity, buffer->host_data, buffer->weight }); + buffer->generation, buffer->capacity, buffer->host_data, buffer->weight, + value->byte_count == 0 }); } return from_bindings(std::move(bindings), result.status); } @@ -42,7 +43,7 @@ CommandProgramBindings CommandProgramBindings::from_bindings(std::vector(binding.offset)); mix_hash(hash, static_cast(binding.length)); mix_hash(hash, binding.weight ? 1 : 0); + mix_hash(hash, binding.empty_value ? 1 : 0); } mix_hash(hash, static_cast(bindings.bindings().size())); return { hash }; @@ -82,7 +84,7 @@ CommandProgramBindingsFingerprint command_program_bindings_fingerprint(const Com out << "|value=" << binding.value.value << "|kind=" << (binding.host_data != nullptr ? "host" : "device") << "|identity=" << binding.identity << "|generation=" << binding.generation << "|capacity=" << binding.capacity << "|offset=" << binding.offset << "|length=" << binding.length - << "|weight=" << (binding.weight ? 1 : 0); + << "|weight=" << (binding.weight ? 1 : 0) << "|empty=" << (binding.empty_value ? 1 : 0); } return { out.str() }; } diff --git a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h index 5f71ab57f985..d43cc5e8f3ee 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-bindings.h +++ b/ggml/src/ggml-hrx/dispatch/command-program-bindings.h @@ -22,6 +22,7 @@ struct CommandProgramBinding { size_t capacity = 0; void * host_data = nullptr; bool weight = false; + bool empty_value = false; bool requires_materialization() const { return host_data != nullptr; } }; diff --git a/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp b/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp index 433204ca7910..8fae265b6e4c 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program-diagnostics.cpp @@ -57,6 +57,11 @@ std::string format_command_binding(const CommandBinding & binding) { << " origin=" << command_binding_origin_name(binding.origin) << " access=" << resource_access_name(binding.access) << " range=[" << binding.offset << ", " << binding.offset + binding.length << ")"; + if (binding.layout != kNativeWeightLayout) { + out << " layout=" << binding.layout << " source_type=" << static_cast(binding.source_type) + << " input_size=" << binding.input_size << " output_size=" << binding.output_size + << " source_length=" << binding.source_length; + } return out.str(); } diff --git a/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp b/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp new file mode 100644 index 000000000000..9f43349a2fff --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-dump.cpp @@ -0,0 +1,365 @@ +#include "command-program-dump.h" + +#include "command-program-diagnostics.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr const char * kDumpDirectoryEnvironment = "GGML_HRX_DUMP_COMMAND_PROGRAM_DIR"; + +static void mix_hash(uint64_t & hash, uint64_t value) { + hash ^= value; + hash *= UINT64_C(1099511628211); +} + +static uint64_t hash_string(const std::string & text) { + uint64_t hash = UINT64_C(1469598103934665603); + for (const char c : text) { + mix_hash(hash, static_cast(c)); + } + return hash; +} + +static std::string hex_u64(uint64_t value) { + std::ostringstream out; + out << std::hex << std::setw(16) << std::setfill('0') << value; + return out.str(); +} + +static std::string dot_escape(const std::string & text) { + std::string escaped; + escaped.reserve(text.size()); + for (const char c : text) { + if (c == '\\' || c == '"') { + escaped.push_back('\\'); + } + if (c == '\n') { + escaped += "\\n"; + } else { + escaped.push_back(c); + } + } + return escaped; +} + +static std::string json_escape(const std::string & text) { + std::string escaped; + escaped.reserve(text.size()); + for (const unsigned char c : text) { + switch (c) { + case '\\': + escaped += "\\\\"; + break; + case '"': + escaped += "\\\""; + break; + case '\b': + escaped += "\\b"; + break; + case '\f': + escaped += "\\f"; + break; + case '\n': + escaped += "\\n"; + break; + case '\r': + escaped += "\\r"; + break; + case '\t': + escaped += "\\t"; + break; + default: + if (c < 0x20) { + static constexpr char kHex[] = "0123456789abcdef"; + escaped += "\\u00"; + escaped.push_back(kHex[c >> 4]); + escaped.push_back(kHex[c & 0xf]); + } else { + escaped.push_back(static_cast(c)); + } + break; + } + } + return escaped; +} + +static std::string kernel_name(const KernelCorpus & corpus, const std::string & target, uint64_t kernel_id) { + const KernelResolveResult resolved = resolve_kernel_definition(corpus, target, kernel_id); + if (resolved.found()) { + return kernel_definition_name(*resolved.definition); + } + std::ostringstream out; + out << "kernel_id=" << kernel_id; + return out.str(); +} + +static void append_json_integer_map(std::ostringstream & out, const std::map & values) { + out << '{'; + size_t index = 0; + for (const auto & value : values) { + if (index++ > 0) { + out << ", "; + } + out << '"' << json_escape(value.first) << "\": " << value.second; + } + out << '}'; +} + +static void append_json_string_map(std::ostringstream & out, const std::map & values) { + out << '{'; + size_t index = 0; + for (const auto & value : values) { + if (index++ > 0) { + out << ", "; + } + out << '"' << json_escape(value.first) << "\": \"" << json_escape(value.second) << '"'; + } + out << '}'; +} + +static void append_kernel_list(std::ostringstream & out, + const std::vector & commands, + const char * phase, + const KernelCorpus & corpus, + const std::string & target) { + for (const Command & command : commands) { + out << phase << ' ' << command.ordinal << ' ' << kernel_name(corpus, target, command.kernel.kernel_id); + if (!command.dependencies.empty()) { + out << " deps="; + for (size_t i = 0; i < command.dependencies.size(); ++i) { + if (i > 0) { + out << ','; + } + out << command.dependencies[i]; + } + } + out << '\n'; + } +} + +static std::string format_kernel_list(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target) { + std::ostringstream out; + append_kernel_list(out, program.initialization_commands, "init", corpus, target); + append_kernel_list(out, program.commands, "main", corpus, target); + return out.str(); +} + +static void append_dot_nodes(std::ostringstream & out, + const std::vector & commands, + const char * prefix, + const char * phase, + const KernelCorpus & corpus, + const std::string & target) { + for (const Command & command : commands) { + const std::string name = kernel_name(corpus, target, command.kernel.kernel_id); + out << " " << prefix << command.ordinal << " [label=\"" << phase << ' ' << command.ordinal << "\\n" + << dot_escape(name) << "\"];\n"; + } +} + +static std::string format_kernel_dot(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target) { + std::ostringstream out; + out << "digraph hrx_kernel_invocations {\n"; + out << " rankdir=LR;\n"; + append_dot_nodes(out, program.initialization_commands, "init", "init", corpus, target); + append_dot_nodes(out, program.commands, "main", "main", corpus, target); + for (const Command & command : program.commands) { + for (const uint32_t dependency : command.dependencies) { + out << " main" << dependency << " -> main" << command.ordinal << ";\n"; + } + } + out << "}\n"; + return out.str(); +} + +static void append_json_command_list(std::ostringstream & out, + const std::vector & commands, + const char * phase, + const KernelCorpus & corpus, + const std::string & target, + const char * indent) { + for (size_t command_index = 0; command_index < commands.size(); ++command_index) { + const Command & command = commands[command_index]; + if (command_index > 0) { + out << ",\n"; + } + const std::string name = kernel_name(corpus, target, command.kernel.kernel_id); + out << indent << "{\n"; + out << indent << " \"phase\": \"" << phase << "\",\n"; + out << indent << " \"ordinal\": " << command.ordinal << ",\n"; + out << indent << " \"kind\": \"" << json_escape(command_kind_name(command.kind)) << "\",\n"; + out << indent << " \"kernel_id\": " << command.kernel.kernel_id << ",\n"; + out << indent << " \"kernel\": \"" << json_escape(name) << "\",\n"; + out << indent << " \"integer_parameters\": "; + append_json_integer_map(out, command.kernel.integer_parameters); + out << ",\n"; + out << indent << " \"compile_parameters\": "; + append_json_string_map(out, command.kernel.compile_parameters); + out << ",\n"; + out << indent << " \"dependencies\": ["; + for (size_t i = 0; i < command.dependencies.size(); ++i) { + if (i > 0) { + out << ", "; + } + out << command.dependencies[i]; + } + out << "],\n"; + out << indent << " \"bindings\": [\n"; + for (size_t binding_index = 0; binding_index < command.bindings.size(); ++binding_index) { + const CommandBinding & binding = command.bindings[binding_index]; + if (binding_index > 0) { + out << ",\n"; + } + out << indent << " {\n"; + out << indent << " \"index\": " << binding_index << ",\n"; + out << indent << " \"name\": \"" << json_escape(binding.name) << "\",\n"; + out << indent << " \"value\": " << binding.value.value << ",\n"; + out << indent << " \"origin\": \"" << json_escape(command_binding_origin_name(binding.origin)) << "\",\n"; + out << indent << " \"access\": \"" << json_escape(resource_access_name(binding.access)) << "\",\n"; + out << indent << " \"offset\": " << binding.offset << ",\n"; + out << indent << " \"length\": " << binding.length << "\n"; + out << indent << " }"; + } + out << '\n' << indent << " ]\n"; + out << indent << '}'; + } +} + +static std::string format_kernel_json(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target, + const std::string & command_shape, + uint64_t dump_id, + const std::string & shape_hash) { + std::ostringstream out; + out << "{\n"; + out << " \"schema\": \"ggml-hrx-command-program-v1\",\n"; + out << " \"dump_id\": " << dump_id << ",\n"; + out << " \"shape_hash\": \"" << json_escape(shape_hash) << "\",\n"; + out << " \"target\": \"" << json_escape(target) << "\",\n"; + out << " \"command_shape\": \"" << json_escape(command_shape) << "\",\n"; + out << " \"transients\": {\n"; + out << " \"arena_size\": " << program.transients.arena_size << ",\n"; + out << " \"arena_alignment\": " << program.transients.arena_alignment << ",\n"; + out << " \"allocations\": [\n"; + for (size_t i = 0; i < program.transients.allocations.size(); ++i) { + const TransientAllocation & allocation = program.transients.allocations[i]; + if (i > 0) { + out << ",\n"; + } + out << " {\"value\": " << allocation.value.value << ", \"size\": " << allocation.size + << ", \"alignment\": " << allocation.alignment << ", \"arena_offset\": " << allocation.arena_offset + << '}'; + } + out << "\n ]\n"; + out << " },\n"; + out << " \"completion_counters\": {\n"; + out << " \"arena_offset\": " << program.completion_counters.arena_offset << ",\n"; + 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()) { + append_json_command_list(out, program.initialization_commands, "init", corpus, target, " "); + wrote_any = true; + } + if (!program.commands.empty()) { + if (wrote_any) { + out << ",\n"; + } + append_json_command_list(out, program.commands, "main", corpus, target, " "); + } + out << "\n ]\n"; + out << "}\n"; + return out.str(); +} + +static void write_file(const std::filesystem::path & path, const std::string & contents) { + std::ofstream output(path, std::ios::binary | std::ios::trunc); + if (!output) { + throw std::runtime_error("cannot create " + path.string()); + } + output << contents; + if (contents.empty() || contents.back() != '\n') { + output << '\n'; + } +} + +} // namespace + +Status dump_command_program_kernels_if_requested(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target, + const std::string & command_shape) { + Status status; + const char * directory_value = std::getenv(kDumpDirectoryEnvironment); + if (directory_value == nullptr || directory_value[0] == '\0' || + (directory_value[0] == '0' && directory_value[1] == '\0')) { + return status; + } + + static std::mutex mutex; + static std::unordered_set dumped_shapes; + static std::atomic sequence{ 0 }; + + const std::filesystem::path root_directory(directory_value); + const std::string dump_key = root_directory.string() + '\n' + target + '\n' + command_shape; + uint64_t dump_id = 0; + { + std::lock_guard lock(mutex); + if (!dumped_shapes.insert(dump_key).second) { + return status; + } + dump_id = sequence.fetch_add(1); + } + + try { + const uint64_t shape_hash = hash_string(dump_key); + const std::string shape_hex = hex_u64(shape_hash); + std::ostringstream name; + name << "program-" << dump_id << "-shape-" << shape_hex; + const std::filesystem::path directory = root_directory / name.str(); + std::filesystem::create_directories(directory); + write_file(directory / "kernels.txt", format_kernel_list(program, corpus, target)); + write_file(directory / "kernels.dot", format_kernel_dot(program, corpus, target)); + write_file(directory / "program.json", format_kernel_json(program, corpus, target, command_shape, dump_id, shape_hex)); + } catch (const std::exception & error) { + status.log("failed to write HRX command program kernel dump: %s", error.what()); + } + return status; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program-dump.h b/ggml/src/ggml-hrx/dispatch/command-program-dump.h new file mode 100644 index 000000000000..f0d911ebdd7f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch/command-program-dump.h @@ -0,0 +1,16 @@ +#pragma once + +#include "command-program.h" +#include "kernel-corpus/kernel-corpus.h" +#include "status.h" + +#include + +namespace ggml::hrx { + +Status dump_command_program_kernels_if_requested(const CommandProgram & program, + const KernelCorpus & corpus, + const std::string & target, + const std::string & command_shape); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch/command-program.cpp b/ggml/src/ggml-hrx/dispatch/command-program.cpp index 642d8589f7f9..4f0127c763a8 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program.cpp +++ b/ggml/src/ggml-hrx/dispatch/command-program.cpp @@ -4,10 +4,12 @@ #include "transient-allocator.h" #include +#include #include #include #include #include +#include #include namespace ggml::hrx { @@ -21,6 +23,11 @@ static std::string string_value(const char * value) { return value != nullptr ? value : ""; } +static bool replay_seeded_transients_enabled() { + const char * value = std::getenv("GGML_HRX_REPLAY_SEEDED_TRANSIENTS"); + return value != nullptr && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0'); +} + static CommandBindingOrigin command_binding_origin(const Graph & graph, const Value & value) { if (value.kind == ValueKind::External) { return CommandBindingOrigin::GraphValue; @@ -40,13 +47,66 @@ struct StorageBindingTarget { size_t offset = 0; }; -static StorageBindingTarget storage_binding_target(const Graph & graph, ValueId value) { +static bool resource_access_writes(ResourceAccess access) { + return access == ResourceAccess::Write || access == ResourceAccess::ReadWrite; +} + +static bool byte_range_is_covered(size_t range_offset, size_t range_length, size_t cover_offset, size_t cover_length) { + if (range_offset < cover_offset) { + return false; + } + const size_t relative_offset = range_offset - cover_offset; + return relative_offset <= cover_length && range_length <= cover_length - relative_offset; +} + +static const Value * find_external_storage_binding_target(const Graph & graph, + const Value & source, + size_t binding_offset, + size_t binding_length) { + if (source.storage.value < 0 || binding_length > source.byte_count || + binding_offset > source.byte_count - binding_length) { + return nullptr; + } + + if (binding_offset > std::numeric_limits::max() - source.storage_offset) { + return nullptr; + } + const size_t range_offset = source.storage_offset + binding_offset; + const Value * best = nullptr; + for (const Value & candidate : graph.values().values()) { + if (candidate.kind != ValueKind::External || candidate.storage != source.storage) { + continue; + } + if (!byte_range_is_covered(range_offset, binding_length, candidate.storage_offset, candidate.byte_count)) { + continue; + } + if (best == nullptr || candidate.byte_count < best->byte_count) { + best = &candidate; + } + } + return best; +} + +static StorageBindingTarget storage_binding_target(const Graph & graph, + ValueId value, + size_t binding_offset, + size_t binding_length, + ResourceAccess access) { StorageBindingTarget target; target.value = value; const Value * graph_value = graph.values().find(value); if (graph_value == nullptr || graph_value->kind != ValueKind::Transient) { return target; } + if (resource_access_writes(access)) { + const Value * external_target = + find_external_storage_binding_target(graph, *graph_value, binding_offset, binding_length); + if (external_target != nullptr) { + target.value = external_target->id; + target.offset = graph_value->storage_offset + binding_offset - external_target->storage_offset; + return target; + } + } const Value * root = graph.values().find(graph_value->storage_root); if (root == nullptr) { return target; @@ -102,6 +162,11 @@ static void append_command(const Graph & graph, command_binding.value = binding.value; command_binding.offset = binding.offset; command_binding.length = binding.length; + command_binding.layout = binding.layout; + command_binding.source_type = binding.source_type; + command_binding.input_size = binding.input_size; + command_binding.output_size = binding.output_size; + command_binding.source_length = binding.source_length; const Value * value = graph.values().find(command_binding.value); const CommandPlanTransient * plan_transient = find_plan_transient(plan, command_binding.value); const CommandPlanCompletionCounterRequest * completion_counter = @@ -109,13 +174,14 @@ static void append_command(const Graph & graph, if (value == nullptr && plan_transient == nullptr && completion_counter == nullptr) { status.log("command %u binding %zu references missing value %d", command.ordinal, binding_index, command_binding.value.value); - } else if (value != nullptr) { - command_binding.origin = command_binding_origin(graph, *value); - } else { - command_binding.origin = CommandBindingOrigin::Transient; } - const StorageBindingTarget binding_target = storage_binding_target(graph, command_binding.value); - command_binding.value = binding_target.value; + if (definition != nullptr && binding_index < definition->bindings.size()) { + command_binding.name = string_value(definition->bindings[binding_index].name); + command_binding.access = definition->bindings[binding_index].access; + } + const StorageBindingTarget binding_target = storage_binding_target( + graph, command_binding.value, command_binding.offset, command_binding.length, command_binding.access); + command_binding.value = binding_target.value; if (binding_target.offset > 0) { if (binding_target.offset > std::numeric_limits::max() - command_binding.offset) { status.log("command %u binding %zu storage alias offset overflows", command.ordinal, binding_index); @@ -123,9 +189,11 @@ static void append_command(const Graph & graph, command_binding.offset += binding_target.offset; } } - if (definition != nullptr && binding_index < definition->bindings.size()) { - command_binding.name = string_value(definition->bindings[binding_index].name); - command_binding.access = definition->bindings[binding_index].access; + const Value * target_value = graph.values().find(command_binding.value); + if (target_value != nullptr) { + command_binding.origin = command_binding_origin(graph, *target_value); + } else if (plan_transient != nullptr || completion_counter != nullptr) { + command_binding.origin = CommandBindingOrigin::Transient; } command.bindings.push_back(std::move(command_binding)); } @@ -140,7 +208,7 @@ static void verify_command_list(const std::vector & commands, for (size_t i = 0; i < commands.size(); ++i) { const Command & command = commands[i]; const std::string command_context = format_command(command); - if (command.ordinal != i) { + if (command.ordinal != i && !replay_seeded_transients_enabled()) { status.log("%s has non-contiguous ordinal at index %zu", command_context.c_str(), i); } if (command.kind != CommandKind::Kernel) { @@ -315,6 +383,78 @@ static void verify_transient_allocations(const CommandProgram & program, Status } } +static bool binding_requires_prior_write(const CommandBinding & binding) { + return binding.origin == CommandBindingOrigin::Transient && binding.access == ResourceAccess::Read; +} + +static bool binding_defines_transient_value(const CommandBinding & binding) { + return binding.origin == CommandBindingOrigin::Transient && resource_access_writes(binding.access); +} + +static int32_t find_later_transient_writer(const std::vector & commands, + size_t begin, + ValueId value) { + for (size_t i = begin; i < commands.size(); ++i) { + const Command & command = commands[i]; + for (const CommandBinding & binding : command.bindings) { + if (binding.value == value && binding_defines_transient_value(binding)) { + return static_cast(command.ordinal); + } + } + } + return -1; +} + +static void verify_transient_reads_are_defined(const CommandProgram & program, Status & status) { + if (replay_seeded_transients_enabled()) { + return; + } + + std::unordered_set defined; + for (const Command & command : program.initialization_commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding_defines_transient_value(binding)) { + defined.insert(binding.value.value); + } + } + } + for (const ConstantInitialization & initialization : program.constant_initializations) { + defined.insert(initialization.value.value); + } + if (program.completion_counters.byte_count > 0) { + for (const TransientAllocation & allocation : program.transients.allocations) { + if (allocation_overlaps_region(allocation, program.completion_counters.arena_offset, + program.completion_counters.byte_count)) { + defined.insert(allocation.value.value); + } + } + } + + for (size_t command_index = 0; command_index < program.commands.size(); ++command_index) { + const Command & command = program.commands[command_index]; + const std::string command_context = format_command(command); + for (size_t binding_index = 0; binding_index < command.bindings.size(); ++binding_index) { + const CommandBinding & binding = command.bindings[binding_index]; + if (!binding_requires_prior_write(binding) || defined.find(binding.value.value) != defined.end()) { + continue; + } + const int32_t later_writer = find_later_transient_writer(program.commands, command_index + 1, binding.value); + if (later_writer >= 0) { + status.log("%s %s reads transient value %d before write by command %d", command_context.c_str(), + format_command_binding(binding).c_str(), binding.value.value, later_writer); + } else { + status.log("%s %s reads transient value %d before write", command_context.c_str(), + format_command_binding(binding).c_str(), binding.value.value); + } + } + for (const CommandBinding & binding : command.bindings) { + if (binding_defines_transient_value(binding)) { + defined.insert(binding.value.value); + } + } + } +} + } // namespace const TransientAllocation * find_transient_allocation(const TransientPlan & plan, ValueId value) { @@ -351,6 +491,7 @@ CommandProgram build_command_program(const Graph & graph, initialization.data, }); } + result.activation_publication_diagnostics = plan.metadata.activation_publication_diagnostics(); return result; } @@ -387,6 +528,7 @@ VerificationResult verify_command_program(const CommandProgram & program, } } verify_transient_allocations(program, result.status); + verify_transient_reads_are_defined(program, result.status); for (const ConstantInitialization & initialization : program.constant_initializations) { const TransientAllocation * allocation = find_transient_allocation(program.transients, initialization.value); if (allocation == nullptr) { diff --git a/ggml/src/ggml-hrx/dispatch/command-program.h b/ggml/src/ggml-hrx/dispatch/command-program.h index 980fe8c3e534..3b7cd8fd7465 100644 --- a/ggml/src/ggml-hrx/dispatch/command-program.h +++ b/ggml/src/ggml-hrx/dispatch/command-program.h @@ -31,6 +31,11 @@ struct CommandBinding { size_t offset = 0; size_t length = 0; ResourceAccess access = ResourceAccess::Read; + std::string layout = kNativeWeightLayout; + ggml_type source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; + size_t source_length = 0; }; struct Command { @@ -73,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 d822691b785d..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); } @@ -115,6 +115,30 @@ static bool can_elide_layout_alias_node(const Graph & graph, return is_layout_alias_node(graph, node) && value_is_available(graph, node.inputs[0], covered_nodes); } +static bool node_inputs_are_available(const Graph & graph, + const GraphNode & node, + const std::vector & covered_nodes) { + for (const ValueId input : node.inputs) { + if (!value_is_available(graph, input, covered_nodes)) { + return false; + } + } + return true; +} + +static bool can_elide_zero_output_node(const Graph & graph, + const GraphNode & node, + const std::vector & covered_nodes) { + const Value * output = graph.values().find(node.output); + return output != nullptr && (output->element_count == 0 || output->byte_count == 0) && + node_inputs_are_available(graph, node, covered_nodes); +} + +static bool can_elide_node(const Graph & graph, const GraphNode & node, const std::vector & covered_nodes) { + return can_elide_layout_alias_node(graph, node, covered_nodes) || + can_elide_zero_output_node(graph, node, covered_nodes); +} + static bool apply_value_aliases(Graph & graph, const DispatchMatch & match, Status & status) { for (const DispatchValueAliasRequest & alias : match.value_aliases) { Status alias_status = graph.values().alias_storage(alias.target_value, alias.source_value); @@ -162,13 +186,17 @@ bool DispatchScheduler::schedule_graph(Graph & graph, if (covered_nodes[i]) { continue; } + if (can_elide_zero_output_node(graph, *node, covered_nodes)) { + covered_nodes[i] = true; + continue; + } DispatchMatch match; const ValueId next_plan_value(static_cast(graph.values().size() + plan_.transients.size() + plan_.completion_counter_requests.size())); DispatchMatchDiagnostics match_diagnostics; if (!try_match_registration(graph, node, i, covered_nodes, plan_, *registry, next_plan_value, match, &match_diagnostics)) { - if (can_elide_layout_alias_node(graph, *node, covered_nodes)) { + if (can_elide_node(graph, *node, covered_nodes)) { pending_diagnostics.append(match.status); covered_nodes[i] = true; continue; @@ -233,7 +261,7 @@ bool DispatchScheduler::schedule_graph(Graph & graph, } for (size_t i = 0; i < nodes.size(); ++i) { if (!covered_nodes[i]) { - if (can_elide_layout_alias_node(graph, nodes[i], covered_nodes)) { + if (can_elide_node(graph, nodes[i], covered_nodes)) { covered_nodes[i] = true; continue; } diff --git a/ggml/src/ggml-hrx/dispatch/dispatch.h b/ggml/src/ggml-hrx/dispatch/dispatch.h index a5e388628800..a59879faef9c 100644 --- a/ggml/src/ggml-hrx/dispatch/dispatch.h +++ b/ggml/src/ggml-hrx/dispatch/dispatch.h @@ -12,8 +12,29 @@ namespace ggml::hrx { +inline constexpr const char kNativeWeightLayout[] = "ggml-native"; +inline constexpr const char kQ4KPackedK256Row64Layout[] = "q4k-packed-k256-row64"; +inline constexpr const char kSymmetricI4K32Row64Layout[] = "symi4-k32-row64"; +inline constexpr const char kSymmetricI2K32EightGroupsShared4Layout[] = "symi2-k32-eightgroups-shared4-payload-first"; +inline constexpr const char kSymmetricI4K32EightGroupsShared4Layout[] = "symi4-k32-eightgroups-shared4-payload-first"; +inline constexpr const char kSymmetricI4K32EightGroupsShared4MultistartLayout[] = + "symi4-k32-eightgroups-shared4-multistart-payload-first"; +inline constexpr const char kQ6KSymmetricI2PackedK256Row64ScaleRowLayout[] = + "q6k-symi2-k32-eightgroups-shared4-plus-packed-k256-row64-scalerow"; +inline constexpr const char kSymmetricI4K64Row64Layout[] = "symi4-k64-row64"; +inline constexpr const char kQ5KSymmetricI5K32Layout[] = "q5k-symi5-k32-native-footprint"; +inline constexpr const char kQ5KSymmetricI8K256Row64Layout[] = "q5k-symi8-k256-row64"; +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; }; @@ -28,6 +49,12 @@ struct DispatchBinding { ValueId value; size_t offset = 0; size_t length = 0; + // Non-native layouts request a backend-resident materialization; the graph value remains unchanged. + std::string layout = kNativeWeightLayout; + ggml_type source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; + size_t source_length = 0; }; struct Dispatch { 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-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary-common.h new file mode 100644 index 000000000000..631747896fe3 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary-common.h @@ -0,0 +1,154 @@ +#pragma once + +#include "../dispatch-registry.h" +#include "graph/op-params.h" + +#include +#include + +namespace ggml::hrx { + +inline constexpr uint32_t common_binary_kind_bit(BinaryKind kind) { + return uint32_t{ 1 } << static_cast(kind); +} + +inline constexpr uint32_t kCommonArithmeticBinaryKinds = + common_binary_kind_bit(BinaryKind::Add) | common_binary_kind_bit(BinaryKind::Sub) | + common_binary_kind_bit(BinaryKind::Mul) | common_binary_kind_bit(BinaryKind::Div); + +inline constexpr uint32_t kCommonGatedBinaryKinds = + common_binary_kind_bit(BinaryKind::SwiGLU) | common_binary_kind_bit(BinaryKind::GeGLU) | + common_binary_kind_bit(BinaryKind::RegLU) | common_binary_kind_bit(BinaryKind::GeGLUErf) | + common_binary_kind_bit(BinaryKind::GeGLUQuick); + +inline bool common_binary_kind_allowed(BinaryKind kind, uint32_t allowed_kinds) { + return binary_kind_supported(kind) && (allowed_kinds & common_binary_kind_bit(kind)) != 0; +} + +inline bool common_binary_kind_requires_order(BinaryKind kind) { + return kind == BinaryKind::Sub || kind == BinaryKind::Div; +} + +inline bool common_binary_kind_matches_graph_op(BinaryKind kind, ggml_op op) { + switch (kind) { + case BinaryKind::Add: + return op == GGML_OP_ADD; + case BinaryKind::Sub: + return op == GGML_OP_SUB; + case BinaryKind::Mul: + return op == GGML_OP_MUL; + case BinaryKind::Div: + return op == GGML_OP_DIV; + case BinaryKind::SwiGLU: + case BinaryKind::GeGLU: + case BinaryKind::RegLU: + case BinaryKind::GeGLUErf: + case BinaryKind::GeGLUQuick: + return op == GGML_OP_GLU; + } + return false; +} + +inline bool common_arithmetic_kind_from_graph_op(ggml_op op, BinaryKind & kind) { + switch (op) { + case GGML_OP_ADD: + kind = BinaryKind::Add; + return true; + case GGML_OP_SUB: + kind = BinaryKind::Sub; + return true; + case GGML_OP_MUL: + kind = BinaryKind::Mul; + return true; + case GGML_OP_DIV: + kind = BinaryKind::Div; + return true; + default: + return false; + } +} + +inline bool common_fused_binary_kind_from_params(const OpParams & params, BinaryKind & kind) { + const BinaryParams * binary_params = op_params_as(params); + if (binary_params != nullptr && binary_kind_supported(binary_params->op)) { + kind = binary_params->op; + return true; + } + + const GluParams * glu_params = op_params_as(params); + if (glu_params == nullptr) { + return false; + } + + switch (glu_params->op) { + case GGML_GLU_OP_REGLU: + kind = BinaryKind::RegLU; + return true; + case GGML_GLU_OP_SWIGLU: + kind = BinaryKind::SwiGLU; + return true; + case GGML_GLU_OP_GEGLU: + kind = BinaryKind::GeGLU; + return true; + case GGML_GLU_OP_GEGLU_ERF: + kind = BinaryKind::GeGLUErf; + return true; + case GGML_GLU_OP_GEGLU_QUICK: + kind = BinaryKind::GeGLUQuick; + return true; + default: + return false; + } +} + +struct CommonBinaryMatch { + const GraphNode * node = nullptr; + ValueId operand = ValueId(-1); + BinaryKind kind = BinaryKind::Add; + bool producer_is_lhs = true; + + bool matched() const { return node != nullptr && operand.value >= 0; } +}; + +inline CommonBinaryMatch common_match_binary_consumer(const GraphNode * node, + ValueId producer, + uint32_t allowed_kinds) { + CommonBinaryMatch match; + if (node == nullptr || node->inputs.size() != 2) { + return match; + } + + BinaryKind kind; + const bool has_kind = + common_fused_binary_kind_from_params(node->params, kind) || + (std::holds_alternative(node->params) && common_arithmetic_kind_from_graph_op(node->op, kind)); + if (!has_kind || !common_binary_kind_matches_graph_op(kind, node->op) || + !common_binary_kind_allowed(kind, allowed_kinds)) { + return match; + } + + const bool producer_is_lhs = node->inputs[0] == producer; + const bool producer_is_rhs = node->inputs[1] == producer; + if (producer_is_lhs == producer_is_rhs) { + return match; + } + + match.node = node; + match.operand = producer_is_lhs ? node->inputs[1] : node->inputs[0]; + match.kind = kind; + match.producer_is_lhs = producer_is_lhs; + return match; +} + +inline void common_set_binary_compile_parameters(KernelSpecialization & kernel, + const char * op_key, + const char * lhs_key, + BinaryKind kind, + bool producer_is_lhs) { + kernel.compile_parameters.emplace(op_key, std::to_string(binary_kind_config_value(kind))); + if (lhs_key != nullptr) { + kernel.compile_parameters.emplace(lhs_key, producer_is_lhs ? "1" : "0"); + } +} + +} // 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 new file mode 100644 index 000000000000..c42ac59a441a --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.cpp @@ -0,0 +1,388 @@ +#include "dispatch-binary.h" + +#include "dispatch-activation-publication.h" +#include "dispatch-layout-utils.h" +#include "dispatch-mul-mat-common.h" +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +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"); + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool positive_shape(const Value & value) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] <= 0) { + return false; + } + } + return value.element_count > 0; +} + +static bool packed_f32_layout(const Value & value) { + size_t expected_stride = sizeof(float); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.nb[i] != expected_stride) { + return false; + } + expected_stride *= static_cast(value.ne[i]); + } + return true; +} + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool storage_ranges_disjoint(const Value & lhs, size_t lhs_byte_count, const Value & rhs, size_t rhs_byte_count) { + if (lhs.storage != rhs.storage) { + return true; + } + if (lhs.storage_offset > std::numeric_limits::max() - lhs_byte_count || + rhs.storage_offset > std::numeric_limits::max() - rhs_byte_count) { + return false; + } + return lhs.storage_offset + lhs_byte_count <= rhs.storage_offset || + rhs.storage_offset + rhs_byte_count <= lhs.storage_offset; +} + +static bool storage_ranges_disjoint(const Value & lhs, const Value & rhs) { + return storage_ranges_disjoint(lhs, lhs.byte_count, rhs, rhs.byte_count); +} + +static bool binary_output_storage_is_safe(const Value & lhs, const Value & rhs, const Value & output) { + return storage_ranges_disjoint(lhs, output) && storage_ranges_disjoint(rhs, output); +} + +static bool binary_noalias_storage_is_safe(const Value & lhs, const Value & rhs, const Value & output) { + return storage_ranges_disjoint(lhs, rhs) && binary_output_storage_is_safe(lhs, rhs, output); +} + +static bool binary_output_storage_is_safe(const Value & lhs, + size_t lhs_byte_count, + const Value & rhs, + size_t rhs_byte_count, + const Value & output) { + return storage_ranges_disjoint(lhs, lhs_byte_count, output, output.byte_count) && + storage_ranges_disjoint(rhs, rhs_byte_count, output, output.byte_count); +} + +static bool supported_source_layout(const Graph & graph, const Value & value) { + if (value.alias_source.value < 0) { + return true; + } + if (value.storage_offset != 0) { + return false; + } + const GraphNode * producer = graph.index().producer(value.id); + return producer != nullptr && producer->op == GGML_OP_RESHAPE; +} + +static bool broadcastable_to(const Value & source, const Value & output) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (source.ne[i] != output.ne[i] && source.ne[i] != 1) { + return false; + } + } + return true; +} + +static bool binary_kind_allows_broadcast(BinaryKind kind, const Value & lhs, const Value & rhs, const Value & output) { + const bool lhs_full = same_shape(lhs, output); + const bool rhs_full = same_shape(rhs, output); + if (!broadcastable_to(lhs, output) || !broadcastable_to(rhs, output) || (!lhs_full && !rhs_full)) { + return false; + } + + switch (kind) { + case BinaryKind::Add: + case BinaryKind::Mul: + return true; + case BinaryKind::Sub: + case BinaryKind::Div: + return lhs_full; + case BinaryKind::SwiGLU: + case BinaryKind::GeGLU: + case BinaryKind::RegLU: + case BinaryKind::GeGLUErf: + case BinaryKind::GeGLUQuick: + return lhs_full && rhs_full; + } + 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; +} + +static void add_binary_shape_parameters(Dispatch & dispatch, + const Value & lhs, + const Value & rhs, + const Value & output) { + dispatch.kernel.integer_parameters.emplace("element_count", output.element_count); + dispatch.kernel.integer_parameters.emplace("ne0", output.ne[0]); + dispatch.kernel.integer_parameters.emplace("ne1", output.ne[1]); + dispatch.kernel.integer_parameters.emplace("ne2", output.ne[2]); + dispatch.kernel.integer_parameters.emplace("ne3", output.ne[3]); + dispatch.kernel.integer_parameters.emplace("src0_element_count", lhs.element_count); + dispatch.kernel.integer_parameters.emplace("src1_element_count", rhs.element_count); +} + +static void add_broadcast_config(Dispatch & dispatch, + const char * config_prefix, + const char * source_prefix, + const Value & source, + const Value & output) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + dispatch.kernel.compile_parameters.emplace( + std::string(config_prefix) + source_prefix + "_broadcast_dim" + std::to_string(i), + std::to_string(broadcast_dim_flag(source, output, i))); + } +} + +static void bind_binary_source_buffer(Dispatch & dispatch, const Value & value, size_t byte_count) { + dispatch.bindings.push_back({ value.storage_root, value.storage_offset, byte_count }); +} + +static void bind_binary_buffers(Dispatch & dispatch, + const Value & lhs, + size_t lhs_byte_count, + const Value & rhs, + size_t rhs_byte_count, + const Value & output) { + bind_binary_source_buffer(dispatch, lhs, lhs_byte_count); + bind_binary_source_buffer(dispatch, rhs, rhs_byte_count); + dispatch.bindings.push_back({ output.id, 0, output.byte_count }); +} + +static void add_binary_strided_parameters(Dispatch & dispatch, + const Value & lhs, + const Value & rhs, + const Value & output, + size_t lhs_byte_count, + size_t rhs_byte_count) { + dispatch.kernel.integer_parameters.emplace("element_count", output.element_count); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(output.ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", std::to_string(output.ne[1])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", std::to_string(output.ne[2])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", + std::to_string(lhs.nb[1] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", + std::to_string(lhs.nb[2] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", + std::to_string(lhs.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", + std::to_string(rhs.nb[1] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", + std::to_string(rhs.nb[2] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride3", + std::to_string(rhs.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_span", + std::to_string(lhs_byte_count / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_span", + std::to_string(rhs_byte_count / sizeof(float))); +} + +static bool match_binary_swiglu_symmetric_i4_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_GLU || node->inputs.size() != 2 || + !common_is_swiglu_params(node->params)) { + return false; + } + + const Value * lhs = graph_value(context.graph, node->inputs[0]); + const Value * rhs = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (lhs == nullptr || rhs == nullptr || output == nullptr || lhs->type != GGML_TYPE_F32 || + rhs->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || !same_shape(*lhs, *output) || + !same_shape(*rhs, *output) || !packed_f32_layout(*lhs) || !packed_f32_layout(*rhs) || + !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) { + return false; + } + + const int64_t input_size = output->ne[0]; + const int64_t token_count = output->element_count / input_size; + if (token_count < 1 || token_count > 16) { + return false; + } + 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; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kBinarySwiGluSymmetricI4K32Kernel); + dispatch.kernel.compile_parameters.emplace("ggml.binary_swiglu_symmetric_i4.input_size", + std::to_string(input_size)); + 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(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)); + return match.status.success(); +} + +static bool match_binary_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->inputs.size() != 2) { + return false; + } + + const BinaryParams * params = op_params_as(node->params); + if (params == nullptr || !binary_kind_supported(params->op)) { + return false; + } + + const Value * output = graph_value(context.graph, node->output); + const Value * lhs = graph_value(context.graph, node->inputs[0]); + const Value * rhs = graph_value(context.graph, node->inputs[1]); + if (output == nullptr || lhs == nullptr || rhs == nullptr) { + return false; + } + + size_t lhs_byte_count = 0; + size_t rhs_byte_count = 0; + if (output->type != GGML_TYPE_F32 || lhs->type != GGML_TYPE_F32 || rhs->type != GGML_TYPE_F32 || + !positive_shape(*output) || !output->contiguous || !packed_f32_layout(*output) || + !strided_f32_storage_span_bytes(*lhs, lhs_byte_count) || + !strided_f32_storage_span_bytes(*rhs, rhs_byte_count) || + output->alias_source.value >= 0 || + !binary_output_storage_is_safe(*lhs, lhs_byte_count, *rhs, rhs_byte_count, *output) || + static_cast(output->element_count) > std::numeric_limits::max()) { + 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_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 { + if (!lhs->contiguous || !rhs->contiguous || !packed_f32_layout(*lhs) || !packed_f32_layout(*rhs)) { + return false; + } + if (!binary_kind_allows_broadcast(params->op, *lhs, *rhs, *output)) { + return false; + } + dispatch.kernel = make_kernel_specialization(publish_q8 ? kBinaryBcF32PublishQ8_1X4Kernel : kBinaryBcF32Kernel); + if (publish_q8) { + dispatch.kernel.integer_parameters.emplace("token_count", output->ne[1]); + 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 match.status.success(); +} + +static void register_binary_dispatch_for(DispatchRegistryBuilder & registry, ggml_op root_op) { + registry.add({ + "common.binary_f32", + root_op, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_binary_f32_dispatch, + }); +} + +} // namespace + +void register_binary_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.binary_swiglu_symmetric_i4_k32", + GGML_OP_GLU, + DispatchMatchKind::Fused, + 100, + DispatchSource::Common, + match_binary_swiglu_symmetric_i4_dispatch, + }); + register_binary_dispatch_for(registry, GGML_OP_ADD); + register_binary_dispatch_for(registry, GGML_OP_SUB); + register_binary_dispatch_for(registry, GGML_OP_MUL); + register_binary_dispatch_for(registry, GGML_OP_DIV); + register_binary_dispatch_for(registry, GGML_OP_GLU); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.h new file mode 100644 index 000000000000..6445daded097 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-binary.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_binary_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp new file mode 100644 index 000000000000..9ccaebe22270 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.cpp @@ -0,0 +1,37 @@ +#include "dispatch-common.h" + +#include "dispatch-binary.h" +#include "dispatch-copy.h" +#include "dispatch-flash-attention.h" +#include "dispatch-gated-mul-mat-id.h" +#include "dispatch-gated-mul-mat.h" +#include "dispatch-gather-add.h" +#include "dispatch-get-rows.h" +#include "dispatch-glu.h" +#include "dispatch-mul-mat-id.h" +#include "dispatch-mul-mat.h" +#include "dispatch-rmsnorm.h" +#include "dispatch-rope-set-rows.h" +#include "dispatch-scale.h" +#include "dispatch-unary.h" + +namespace ggml::hrx { + +void register_common_dispatches(DispatchRegistryBuilder & registry) { + register_binary_dispatch(registry); + register_copy_dispatch(registry); + register_flash_attention_dispatches(registry); + register_gated_mul_mat_id_dispatches(registry); + register_gated_mul_mat_dispatches(registry); + register_gather_add_dispatch(registry); + register_get_rows_dispatches(registry); + register_glu_dispatches(registry); + register_mul_mat_id_dispatches(registry); + register_mul_mat_dispatches(registry); + register_rope_set_rows_dispatches(registry); + register_scale_dispatch(registry); + register_unary_dispatch(registry); + register_rmsnorm_dispatches(registry); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h new file mode 100644 index 000000000000..d7fe2aed8246 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-common.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_common_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp new file mode 100644 index 000000000000..5f091f40e564 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.cpp @@ -0,0 +1,231 @@ +#include "dispatch-copy.h" + +#include "dispatch-layout-utils.h" +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr uint64_t kMaximumCopyElements = uint64_t{ 1 } << 30; +static constexpr KernelCatalogRef kCopyF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_f32"); +static constexpr KernelCatalogRef kCopyStridedSourceF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_strided_source_f32"); +static constexpr KernelCatalogRef kConcatDim0F32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_concat_dim0_f32"); +static constexpr KernelCatalogRef kConcatDim0StridedSourceF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_concat_dim0_strided_source_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool packed_f32_layout(const Value & value) { + size_t expected_stride = sizeof(float); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.nb[i] != expected_stride) { + return false; + } + expected_stride *= static_cast(value.ne[i]); + } + return true; +} + +static bool f32_element_strides(const Value & value) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] <= 0 || value.nb[i] <= 0 || value.nb[i] % sizeof(float) != 0) { + return false; + } + } + return value.byte_count % sizeof(float) == 0; +} + +static void add_concat_strided_source_parameters(KernelSpecialization & kernel, const Value & lhs, const Value & rhs) { + static constexpr const char * kPrefix = "ggml.concat_dim0_strided_source_f32."; + kernel.integer_parameters.emplace("lhs_span", lhs.byte_count / sizeof(float)); + kernel.integer_parameters.emplace("rhs_span", rhs.byte_count / sizeof(float)); + kernel.compile_parameters.emplace(std::string(kPrefix) + "lhs_width", std::to_string(lhs.ne[0])); + kernel.compile_parameters.emplace(std::string(kPrefix) + "rhs_width", std::to_string(rhs.ne[0])); + kernel.compile_parameters.emplace(std::string(kPrefix) + "row_count", + std::to_string(lhs.ne[1] * lhs.ne[2] * lhs.ne[3])); + kernel.compile_parameters.emplace(std::string(kPrefix) + "row_ne1", std::to_string(lhs.ne[1])); + kernel.compile_parameters.emplace(std::string(kPrefix) + "row_ne2", std::to_string(lhs.ne[2])); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + kernel.compile_parameters.emplace(std::string(kPrefix) + "lhs_stride" + std::to_string(i), + std::to_string(lhs.nb[i] / sizeof(float))); + kernel.compile_parameters.emplace(std::string(kPrefix) + "rhs_stride" + std::to_string(i), + std::to_string(rhs.nb[i] / sizeof(float))); + } +} + +static bool make_copy_f32_dispatch(const Value & source, const Value & output, Dispatch & dispatch) { + if (source.type != GGML_TYPE_F32 || output.type != GGML_TYPE_F32 || !output.contiguous || + !packed_f32_layout(output) || !f32_element_strides(source) || source.element_count <= 0 || + source.element_count != output.element_count || + static_cast(source.element_count) > kMaximumCopyElements || source.storage == output.storage) { + return false; + } + + 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; + } + const size_t source_span_elements = source_span / sizeof(float); + if (source_span_elements > kMaximumCopyElements) { + return false; + } + dispatch.kernel = make_kernel_specialization(kCopyStridedSourceF32Kernel); + dispatch.kernel.integer_parameters.emplace("source_span", source_span_elements); + dispatch.kernel.compile_parameters.emplace("ggml.copy_strided_source_f32.ne0", std::to_string(source.ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.copy_strided_source_f32.ne1", std::to_string(source.ne[1])); + dispatch.kernel.compile_parameters.emplace("ggml.copy_strided_source_f32.ne2", std::to_string(source.ne[2])); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + dispatch.kernel.compile_parameters.emplace("ggml.copy_strided_source_f32.stride" + std::to_string(i), + std::to_string(source.nb[i] / sizeof(float))); + } + } + dispatch.kernel.integer_parameters.emplace("element_count", source.element_count); + dispatch.bindings.push_back({ source.storage_root, source.storage_offset, source_span }); + dispatch.bindings.push_back({ output.id, 0, output.byte_count }); + return true; +} + +static bool match_copy_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_CPY || node->inputs.size() != 2) { + return false; + } + + const Value * source = graph_value(context.graph, node->inputs[0]); + const Value * target = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (source == nullptr || target == nullptr || output == nullptr || target->type != GGML_TYPE_F32 || + !target->contiguous || !packed_f32_layout(*target) || target->element_count != output->element_count || + target->byte_count != output->byte_count || source->storage == target->storage) { + return false; + } + + Dispatch dispatch; + if (!make_copy_f32_dispatch(*source, *output, dispatch)) { + return false; + } + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_cont_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_CONT || node->inputs.size() != 1) { + return false; + } + + const Value * source = graph_value(context.graph, node->inputs[0]); + const Value * output = graph_value(context.graph, node->output); + if (source == nullptr || output == nullptr) { + return false; + } + + Dispatch dispatch; + if (!make_copy_f32_dispatch(*source, *output, dispatch)) { + return false; + } + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_concat_dim0_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_CONCAT || node->inputs.size() != 2) { + return false; + } + + const Value * lhs = graph_value(context.graph, node->inputs[0]); + const Value * rhs = graph_value(context.graph, node->inputs[1]); + const Value * output = graph_value(context.graph, node->output); + if (lhs == nullptr || rhs == nullptr || output == nullptr || lhs->type != GGML_TYPE_F32 || + rhs->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || !output->contiguous || + !packed_f32_layout(*output) || !f32_element_strides(*lhs) || !f32_element_strides(*rhs) || lhs->ne[0] <= 0 || + lhs->ne[0] > 65536 || rhs->ne[0] <= 0 || rhs->ne[0] > 65536 || output->ne[0] != lhs->ne[0] + rhs->ne[0] || + output->element_count <= 0 || static_cast(output->element_count) > kMaximumCopyElements || + lhs->storage == rhs->storage || lhs->storage == output->storage || rhs->storage == output->storage) { + return false; + } + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + if (lhs->ne[dim] != rhs->ne[dim] || lhs->ne[dim] != output->ne[dim]) { + return false; + } + } + + const int64_t row_count = output->element_count / output->ne[0]; + if (row_count < 1 || row_count > 1048576) { + return false; + } + + 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])); + } else { + const size_t lhs_span = lhs->byte_count / sizeof(float); + const size_t rhs_span = rhs->byte_count / sizeof(float); + if (lhs_span > kMaximumCopyElements || rhs_span > kMaximumCopyElements) { + return false; + } + dispatch.kernel = make_kernel_specialization(kConcatDim0StridedSourceF32Kernel); + add_concat_strided_source_parameters(dispatch.kernel, *lhs, *rhs); + } + dispatch.bindings.push_back({ lhs->id, 0, lhs->byte_count }); + dispatch.bindings.push_back({ rhs->id, 0, rhs->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +} // namespace + +void register_copy_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.copy_f32", + GGML_OP_CPY, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_copy_f32_dispatch, + }); + registry.add({ + "common.cont_f32", + GGML_OP_CONT, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_cont_f32_dispatch, + }); + registry.add({ + "common.concat_dim0_f32", + GGML_OP_CONCAT, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_concat_dim0_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h new file mode 100644 index 000000000000..6228aac1946b --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-copy.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch_registration/dispatch-registry.h" + +namespace ggml::hrx { + +void register_copy_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp new file mode 100644 index 000000000000..c1fec1a8a8ce --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.cpp @@ -0,0 +1,708 @@ +#include "dispatch-flash-attention.h" + +#include "dispatch-activation-publication.h" +#include "dispatch-mul-mat-common.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kFlashAttentionF32F16WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_f32_f16_wmma"); +static constexpr KernelCatalogRef kFlashAttentionF32F16WmmaPublishF16Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_f32_f16_wmma_publish_f16"); +static constexpr KernelCatalogRef kFlashAttentionDecodeSplitNextQ8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); +static constexpr KernelCatalogRef kCopyTransposeF16Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_transpose_f16"); +static constexpr int64_t kDecodeRowCapacity = 16; +static constexpr int64_t kDecodeKvTileSize = 64; +static constexpr int64_t kPrefillQkHeadSizeBlock = 16; +static constexpr int64_t kPrefillValueHeadSizeBlock = 64; +static constexpr int64_t kPrefillMinQkHeadSize = kPrefillQkHeadSizeBlock; +static constexpr int64_t kPrefillMinValueHeadSize = kPrefillValueHeadSizeBlock; +// Current cap keeps full-head Q/K staging within the kernel's fixed LDS budget. +static constexpr int64_t kPrefillMaxHeadSize = 512; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +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; +} + +static bool is_supported_key_value_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 32768; +} + +static bool is_supported_decode_key_value_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +static bool is_supported_decode_query_length(int64_t query_length) { + return query_length >= 1 && query_length < kDecodeRowCapacity; +} + +static bool is_supported_head_count(int64_t head_count) { + return head_count >= 1 && head_count <= 64; +} + +static bool is_supported_qk_head_size(int64_t head_size) { + return head_size >= kPrefillMinQkHeadSize && head_size <= kPrefillMaxHeadSize && + head_size % kPrefillQkHeadSizeBlock == 0; +} + +static bool is_supported_value_head_size(int64_t head_size) { + return head_size >= kPrefillMinValueHeadSize && head_size <= kPrefillMaxHeadSize && + head_size % kPrefillValueHeadSizeBlock == 0; +} + +static bool has_query_layout(const Value & value, int64_t query_head_count, int64_t head_size) { + const size_t element_size = sizeof(float); + return value.nb[0] == element_size && + value.nb[1] == static_cast(query_head_count * head_size) * element_size && + (value.ne[2] == 1 || value.nb[2] == static_cast(head_size) * element_size); +} + +static bool has_key_value_layout(const Value & value, int64_t key_value_head_count, int64_t head_size) { + const size_t element_size = sizeof(ggml_fp16_t); + return value.nb[0] == element_size && + value.nb[1] == static_cast(key_value_head_count * head_size) * element_size && + (value.ne[2] == 1 || value.nb[2] == static_cast(head_size) * element_size); +} + +static bool has_mask_layout(const Value & value, int64_t key_value_token_count) { + const size_t element_size = sizeof(ggml_fp16_t); + return value.nb[0] == element_size && value.nb[1] == static_cast(key_value_token_count) * element_size; +} + +static bool has_output_layout(const Value & value, int64_t query_head_count, int64_t head_size) { + const size_t element_size = sizeof(float); + return value.nb[0] == element_size && value.nb[1] == static_cast(head_size) * element_size && + value.nb[2] == static_cast(query_head_count * head_size) * element_size; +} + +static bool has_flash_attention_params(const GraphNode & node) { + const FlashAttnExtParams * params = op_params_as(node.params); + if (params == nullptr) { + return false; + } + return nearly_equal(params->max_bias, 0.0f) && nearly_equal(params->logit_softcap, 0.0f) && + (params->prec == GGML_PREC_DEFAULT || params->prec == GGML_PREC_F32); +} + +static size_t attention_mask_byte_count(int64_t query_token_count, int64_t key_value_token_count) { + if (query_token_count <= 0 || key_value_token_count <= 0) { + return 0; + } + return static_cast(query_token_count) * static_cast(key_value_token_count) * sizeof(ggml_fp16_t); +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +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); +} + +struct FlashAttentionMatch { + const Value * query = nullptr; + const Value * key = nullptr; + const Value * value = nullptr; + const Value * mask = nullptr; + const Value * output = nullptr; + const GraphNode * output_layout = nullptr; + ValueId mask_binding_value; + size_t mask_binding_bytes = 0; + int64_t query_token_count = 0; + int64_t key_value_token_count = 0; + int64_t query_head_count = 0; + int64_t key_value_head_count = 0; + int64_t qk_head_size = 0; + int64_t value_head_size = 0; + float attention_scale = 0.0f; + + bool matched() const { + return query != nullptr && key != nullptr && value != nullptr && mask != nullptr && output != nullptr; + } +}; + +struct DecodeSplitFlashAttentionMatch { + const Value * query = nullptr; + const Value * key = nullptr; + const Value * value = nullptr; + const Value * mask = nullptr; + const Value * output = nullptr; + const GraphNode * output_layout = nullptr; + int64_t query_token_count = 0; + int64_t key_value_token_count = 0; + int64_t key_value_capacity = 0; + int64_t query_head_count = 0; + int64_t key_value_head_count = 0; + int64_t qk_head_size = 0; + int64_t value_head_size = 0; + float attention_scale = 0.0f; + + bool matched() const { + return query != nullptr && key != nullptr && value != nullptr && mask != nullptr && output != nullptr; + } +}; + +static FlashAttentionMatch match_flash_attention_f32_f16(const Graph & graph, + const CommandPlan & plan, + const GraphNode * node) { + FlashAttentionMatch match; + if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4) { + return match; + } + + const Value * query = graph_value(graph, node->inputs[0]); + const Value * key = graph_value(graph, node->inputs[1]); + const Value * value = graph_value(graph, node->inputs[2]); + const Value * mask = graph_value(graph, node->inputs[3]); + const Value * output = graph_value(graph, node->output); + if (query == nullptr || key == nullptr || value == nullptr || mask == nullptr || output == nullptr) { + return {}; + } + if (query->type != GGML_TYPE_F32 || key->type != GGML_TYPE_F16 || value->type != GGML_TYPE_F16 || + mask->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F32) { + return {}; + } + const int64_t qk_head_size = query->ne[0]; + const int64_t value_head_size = value->ne[0]; + if (!is_supported_qk_head_size(qk_head_size) || !is_supported_value_head_size(value_head_size) || + !has_flash_attention_params(*node)) { + return {}; + } + if (key->ne[0] != qk_head_size || output->ne[0] != value_head_size) { + return {}; + } + if (query->ne[3] != 1 || key->ne[3] != 1 || value->ne[3] != 1 || output->ne[3] != 1 || mask->ne[2] != 1 || + mask->ne[3] != 1) { + return {}; + } + + const int64_t query_token_count = query->ne[1]; + const int64_t query_head_count = query->ne[2]; + const int64_t key_value_capacity = key->ne[1]; + const int64_t key_value_head_count = key->ne[2]; + if (!is_supported_token_count(query_token_count) || key_value_capacity < query_token_count || + !is_supported_head_count(query_head_count) || !is_supported_head_count(key_value_head_count) || + query_head_count % key_value_head_count != 0) { + return {}; + } + if (value->ne[1] != key_value_capacity || value->ne[2] != key_value_head_count || + mask->ne[0] > key_value_capacity || mask->ne[1] != query_token_count || output->ne[1] != query_head_count || + output->ne[2] != query_token_count) { + return {}; + } + + int64_t key_value_token_count = key_value_capacity; + ValueId mask_binding_value = mask->id; + size_t mask_binding_bytes = mask->byte_count; + bool mask_binding_is_alternate = false; + const size_t compact_mask_bytes = attention_mask_byte_count(query_token_count, query_token_count); + const auto * compact_mask = find_alternate_value(plan, mask->id, GGML_TYPE_F16, compact_mask_bytes); + const bool mask_is_compact = mask->ne[0] == query_token_count; + const bool mask_is_capacity = mask->ne[0] == key_value_capacity; + if (compact_mask != nullptr && mask_is_capacity && mask->ne[0] > query_token_count) { + key_value_token_count = query_token_count; + mask_binding_value = compact_mask->alternate_value; + mask_binding_bytes = compact_mask->byte_count; + mask_binding_is_alternate = true; + } else if (!mask_is_compact && !mask_is_capacity) { + return {}; + } + if (!is_supported_key_value_token_count(key_value_token_count)) { + return {}; + } + + if (!has_query_layout(*query, query_head_count, qk_head_size) || + !has_key_value_layout(*key, key_value_head_count, qk_head_size) || + !has_key_value_layout(*value, key_value_head_count, value_head_size) || + !has_output_layout(*output, query_head_count, value_head_size)) { + return {}; + } + if (!mask_binding_is_alternate && !has_mask_layout(*mask, key_value_token_count)) { + return {}; + } + + match.query = query; + match.key = key; + match.value = value; + match.mask = mask; + match.output = output; + match.output_layout = find_single_layout_alias_consumer(graph, output->id); + match.mask_binding_value = mask_binding_value; + match.mask_binding_bytes = mask_binding_bytes; + match.query_token_count = query_token_count; + match.key_value_token_count = 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; + match.value_head_size = value_head_size; + match.attention_scale = op_params_as(node->params)->scale; + return match; +} + +static DecodeSplitFlashAttentionMatch match_decode_split_flash_attention_f32_f16(const Graph & graph, + const GraphNode * node) { + DecodeSplitFlashAttentionMatch match; + if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4) { + return match; + } + + const Value * query = graph_value(graph, node->inputs[0]); + const Value * key = graph_value(graph, node->inputs[1]); + const Value * value = graph_value(graph, node->inputs[2]); + const Value * mask = graph_value(graph, node->inputs[3]); + const Value * output = graph_value(graph, node->output); + if (query == nullptr || key == nullptr || value == nullptr || mask == nullptr || output == nullptr) { + return {}; + } + if (query->type != GGML_TYPE_F32 || key->type != GGML_TYPE_F16 || value->type != GGML_TYPE_F16 || + mask->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F32) { + return {}; + } + + const int64_t qk_head_size = query->ne[0]; + const int64_t value_head_size = value->ne[0]; + if (!is_supported_qk_head_size(qk_head_size) || !is_supported_value_head_size(value_head_size) || + !has_flash_attention_params(*node)) { + return {}; + } + if (key->ne[0] != qk_head_size || output->ne[0] != value_head_size) { + return {}; + } + if (query->ne[3] != 1 || key->ne[3] != 1 || value->ne[3] != 1 || output->ne[3] != 1 || mask->ne[2] != 1 || + mask->ne[3] != 1) { + return {}; + } + + const int64_t query_token_count = query->ne[1]; + const int64_t query_head_count = query->ne[2]; + const int64_t key_value_capacity = key->ne[1]; + const int64_t key_value_head_count = key->ne[2]; + const int64_t key_value_token_count = mask->ne[0]; + if (!is_supported_decode_query_length(query_token_count) || + !is_supported_decode_key_value_token_count(key_value_token_count) || + key_value_capacity < key_value_token_count || !is_supported_head_count(query_head_count) || + !is_supported_head_count(key_value_head_count) || query_head_count % key_value_head_count != 0) { + return {}; + } + if (value->ne[1] != key_value_capacity || value->ne[2] != key_value_head_count || + mask->ne[1] != query_token_count || output->ne[1] != query_head_count || output->ne[2] != query_token_count) { + return {}; + } + if (!has_query_layout(*query, query_head_count, qk_head_size) || + !has_key_value_layout(*key, key_value_head_count, qk_head_size) || + !has_key_value_layout(*value, key_value_head_count, value_head_size) || + !has_mask_layout(*mask, key_value_token_count) || + !has_output_layout(*output, query_head_count, value_head_size)) { + return {}; + } + + match.query = query; + match.key = key; + match.value = value; + match.mask = mask; + match.output = output; + 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 = 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; + match.value_head_size = value_head_size; + match.attention_scale = op_params_as(node->params)->scale; + return match; +} + +static std::string to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +static void add_flash_attention_decode_compile_parameters(KernelSpecialization & kernel, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t qk_head_size, + int64_t value_head_size, + float attention_scale) { + kernel.compile_parameters.emplace("ggml.flash_attention.query_head_count", to_config_value(query_head_count)); + kernel.compile_parameters.emplace("ggml.flash_attention.key_value_head_count", + to_config_value(key_value_head_count)); + kernel.compile_parameters.emplace("ggml.flash_attention.qk_head_size", to_config_value(qk_head_size)); + kernel.compile_parameters.emplace("ggml.flash_attention.value_head_size", to_config_value(value_head_size)); + kernel.compile_parameters.emplace("ggml.flash_attention.attention_scale", to_config_value(attention_scale)); +} + +static void add_flash_attention_compile_parameters(KernelSpecialization & kernel, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t qk_head_size, + int64_t value_head_size, + float attention_scale, + bool apply_gate, + int64_t gate_stride_head, + int64_t gate_stride_token) { + kernel.compile_parameters.emplace("ggml.flash_attention.query_head_count", to_config_value(query_head_count)); + kernel.compile_parameters.emplace("ggml.flash_attention.key_value_head_count", + to_config_value(key_value_head_count)); + kernel.compile_parameters.emplace("ggml.flash_attention.qk_head_size", to_config_value(qk_head_size)); + kernel.compile_parameters.emplace("ggml.flash_attention.value_head_size", to_config_value(value_head_size)); + kernel.compile_parameters.emplace("ggml.flash_attention.attention_scale", to_config_value(attention_scale)); + kernel.compile_parameters.emplace("ggml.flash_attention.apply_gate", apply_gate ? "1" : "0"); + kernel.compile_parameters.emplace("ggml.flash_attention.gate_stride_head", to_config_value(gate_stride_head)); + kernel.compile_parameters.emplace("ggml.flash_attention.gate_stride_token", to_config_value(gate_stride_token)); +} + +static const GraphNode * single_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const std::vector & consumers = graph.index().consumers(value); + return consumers.size() == 1 && consumers.front() != nullptr && consumers.front()->op == op ? consumers.front() : + nullptr; +} + +static DispatchBinding prepare_flash_attention_value(const DispatchMatchContext & context, + const FlashAttentionMatch & match, + DispatchMatch & dispatch_match, + KernelSpecialization & attention) { + if (match.query_token_count < 256 || match.key_value_token_count < 512 || + match.key_value_token_count % 32 != 0) { + return { match.value->id, 0, match.value->byte_count }; + } + + const int64_t columns = match.key_value_head_count * match.value_head_size; + const size_t bytes = static_cast(match.key_value_token_count * columns) * sizeof(ggml_fp16_t); + const ValueId transposed = match_value(context, dispatch_match, 0); + dispatch_match.transients.push_back({ transposed, "common.flash_attention.transposed_value", bytes, 256 }); + + Dispatch copy; + copy.kernel = make_kernel_specialization(kCopyTransposeF16Kernel); + 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)); + attention.compile_parameters.emplace("ggml.flash_attention.value_layout", "1"); + return { transposed, 0, bytes }; +} + +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) { + return false; + } + + const GraphNode * reshape = match.output_layout; + const Value * reshaped = graph_value(context.graph, reshape->output); + const GraphNode * mul = + reshaped != nullptr ? single_consumer_with_op(context.graph, reshaped->id, GGML_OP_MUL) : nullptr; + if (reshape->inputs.size() != 1 || reshaped == nullptr || mul == nullptr || mul->inputs.size() != 2 || + reshaped->type != GGML_TYPE_F32 || !reshaped->contiguous || + reshaped->ne[0] != match.value_head_size * match.query_head_count || + reshaped->ne[1] != match.query_token_count || reshaped->ne[2] != 1 || reshaped->ne[3] != 1) { + return false; + } + + ValueId gate_value_id; + if (mul->inputs[0] == reshape->output) { + gate_value_id = mul->inputs[1]; + } else if (mul->inputs[1] == reshape->output) { + gate_value_id = mul->inputs[0]; + } else { + return false; + } + + const GraphNode * sigmoid = context.graph.index().producer(gate_value_id); + const UnaryParams * sigmoid_params = sigmoid != nullptr ? op_params_as(sigmoid->params) : nullptr; + if (sigmoid == nullptr || sigmoid->op != GGML_OP_UNARY || sigmoid->inputs.size() != 1 || + sigmoid_params == nullptr || sigmoid_params->op != UnaryKind::Sigmoid || + single_consumer_with_op(context.graph, sigmoid->output, GGML_OP_MUL) != mul) { + return false; + } + + const GraphNode * cont = context.graph.index().producer(sigmoid->inputs[0]); + if (cont == nullptr || cont->op != GGML_OP_CONT || cont->inputs.size() != 1 || + single_consumer_with_op(context.graph, cont->output, GGML_OP_UNARY) != sigmoid) { + return false; + } + + const GraphNode * gate_view = context.graph.index().producer(cont->inputs[0]); + const Value * raw_gate = gate_view != nullptr ? graph_value(context.graph, gate_view->output) : nullptr; + const Value * output = graph_value(context.graph, mul->output); + if (gate_view == nullptr || gate_view->op != GGML_OP_VIEW || gate_view->inputs.size() != 1 || raw_gate == nullptr || + output == nullptr || raw_gate->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + raw_gate->ne[0] != match.value_head_size || raw_gate->ne[1] != match.query_head_count || + raw_gate->ne[2] != match.query_token_count || raw_gate->ne[3] != 1 || raw_gate->nb[0] != sizeof(float) || + output->ne != reshaped->ne || !output->contiguous || + single_consumer_with_op(context.graph, gate_view->output, GGML_OP_CONT) != cont) { + return false; + } + if (output->storage_root == match.query->storage_root || output->storage_root == match.key->storage_root || + output->storage_root == match.value->storage_root || output->storage_root == match.mask->storage_root || + output->storage_root == raw_gate->storage_root) { + return false; + } + + dispatch_match.covered_nodes.push_back(context.root_index); + for (const GraphNode * covered : { reshape, gate_view, cont, sigmoid, mul }) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, covered, + dispatch_match.covered_nodes)) { + return false; + } + } + + 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, + match.qk_head_size, match.value_head_size, match.attention_scale, true, + static_cast(raw_gate->nb[1] / sizeof(float)), + static_cast(raw_gate->nb[2] / sizeof(float))); + dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); + dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); + dispatch.bindings.push_back(prepare_flash_attention_value(context, match, dispatch_match, dispatch.kernel)); + dispatch.bindings.push_back({ match.mask_binding_value, 0, match.mask_binding_bytes }); + dispatch.bindings.push_back({ raw_gate->id, 0, raw_gate->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_flash_attention_f32_f16_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const FlashAttentionMatch match = match_flash_attention_f32_f16(context.graph, context.plan, context.root_node); + if (!match.matched()) { + return false; + } + + 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, + match.qk_head_size, match.value_head_size, match.attention_scale, false, 1, + 1); + dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); + dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); + dispatch.bindings.push_back(prepare_flash_attention_value(context, match, dispatch_match, dispatch.kernel)); + dispatch.bindings.push_back({ match.mask_binding_value, 0, match.mask_binding_bytes }); + dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + if (publish_f16) { + constexpr const char * name = "common.flash_attention.f16"; + 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(publication.binding()); + } + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.output_layout != nullptr) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.output_layout, + dispatch_match.covered_nodes)) { + return false; + } + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_flash_attention_decode_split_next_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const DecodeSplitFlashAttentionMatch match = + match_decode_split_flash_attention_f32_f16(context.graph, context.root_node); + if (!match.matched()) { + 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) * + static_cast(kDecodeRowCapacity); + const size_t partial_value_count = partial_scalar_count * static_cast(match.value_head_size); + const size_t partial_scalar_bytes = partial_scalar_count * sizeof(float); + 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; + if (partial_scalar_bytes == 0 || partial_output_bytes == 0) { + return false; + } + + const ValueId partial_max = match_value(context, dispatch_match, 0); + 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); + + dispatch_match.transients.push_back( + { partial_max, "common.decode.flash_attention.partial_max", partial_scalar_bytes, 256 }); + dispatch_match.transients.push_back( + { 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.completion_counter_requests.push_back({ + completion_counter, + "common.decode.flash_attention.completion_counter", + static_cast(match.key_value_head_count), + }); + + 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); + const size_t output_row_bytes = static_cast(output_hidden_size) * sizeof(float); + 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, + match.value_head_size, match.attention_scale); + dispatch.kernel.compile_parameters.emplace("ggml.flash_attention.decode.key_value_token_capacity", + to_config_value(match.key_value_capacity)); + dispatch.bindings.push_back( + { match.query->id, static_cast(row) * match.query->nb[1], query_row_bytes }); + dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); + dispatch.bindings.push_back({ match.value->id, 0, match.value->byte_count }); + dispatch.bindings.push_back({ match.mask->id, static_cast(row) * match.mask->nb[1], mask_row_bytes }); + dispatch.bindings.push_back({ partial_max, 0, partial_scalar_bytes }); + dispatch.bindings.push_back({ partial_sum, 0, partial_scalar_bytes }); + dispatch.bindings.push_back({ partial_output, 0, partial_output_bytes }); + dispatch.bindings.push_back( + { 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(publication.binding(static_cast(row) * q8_row_bytes, q8_row_bytes)); + dispatch_match.dispatches.push_back(std::move(dispatch)); + } + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.output_layout != nullptr) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.output_layout, + dispatch_match.covered_nodes)) { + return false; + } + } + return true; +} + +} // namespace + +void register_flash_attention_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.flash_attention_f32_f16_gate", + GGML_OP_FLASH_ATTN_EXT, + DispatchMatchKind::Fused, + 100, + DispatchSource::Common, + match_flash_attention_gate_dispatch, + }); + registry.add({ + "common.flash_attention_decode_split_next_q8", + GGML_OP_FLASH_ATTN_EXT, + DispatchMatchKind::SingleOp, + 75, + DispatchSource::Common, + match_flash_attention_decode_split_next_q8_dispatch, + }); + registry.add({ + "common.flash_attention_f32_f16_wmma", + GGML_OP_FLASH_ATTN_EXT, + DispatchMatchKind::SingleOp, + 50, + DispatchSource::Common, + match_flash_attention_f32_f16_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.h new file mode 100644 index 000000000000..125a0beeffce --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-flash-attention.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_flash_attention_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp new file mode 100644 index 000000000000..901ddeb9e9b1 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.cpp @@ -0,0 +1,220 @@ +#include "dispatch-gated-mul-mat-id.h" + +#include "dispatch-mul-mat-id-common.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kMulMatIdTiledSwiGLUF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnySwiGLUF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32"); + +struct MulMatIdSwiGLUMatch { + const Value * input = nullptr; + const Value * route_ids = nullptr; + const Value * gate_weight = nullptr; + const Value * up_weight = nullptr; + const Value * gate_output = nullptr; + const Value * up_output = nullptr; + const Value * output = nullptr; + const GraphNode * gate_node = nullptr; + const GraphNode * up_node = nullptr; + const GraphNode * glu_node = nullptr; + CommandPlanMoeRoutingBundle routing_bundle; + bool has_routing_bundle = false; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + int64_t route_count = 0; + int64_t input_route_count = 0; + int64_t expert_count = 0; + CommonMulMatWeightFormat gate_format = CommonMulMatWeightFormat::Q4K; + CommonMulMatWeightFormat up_format = CommonMulMatWeightFormat::Q4K; + + bool matched() const { + return input != nullptr && route_ids != nullptr && gate_weight != nullptr && up_weight != nullptr && + gate_output != nullptr && up_output != nullptr && output != nullptr && gate_node != nullptr && + up_node != nullptr && glu_node != nullptr; + } +}; + +static MulMatIdSwiGLUMatch match_mul_mat_id_swiglu(const DispatchMatchContext & context) { + MulMatIdSwiGLUMatch match; + const CommonMulMatIdMatch root = common_match_mul_mat_id_any_format(context.graph, context.root_node, context.plan); + if (!root.matched() || !context.graph.has_index()) { + return match; + } + + const std::vector & root_consumers = context.graph.index().consumers(context.root_node->output); + if (root_consumers.size() != 1 || root_consumers.front() == nullptr || root_consumers.front()->op != GGML_OP_GLU) { + return {}; + } + + const GraphNode * glu_node = root_consumers.front(); + if (glu_node->inputs.size() != 2 || !common_mul_mat_id_is_swiglu_params(glu_node->params)) { + return {}; + } + + size_t glu_index = 0; + if (!context.graph.index().node_index(glu_node, glu_index) || glu_index >= context.covered_nodes.size() || + context.covered_nodes[glu_index]) { + return {}; + } + + const bool root_is_gate = glu_node->inputs[0] == context.root_node->output; + const bool root_is_up = glu_node->inputs[1] == context.root_node->output; + if (!root_is_gate && !root_is_up) { + return {}; + } + + const ValueId peer_output_id = root_is_gate ? glu_node->inputs[1] : glu_node->inputs[0]; + const GraphNode * peer_node = context.graph.index().producer(peer_output_id); + if (peer_node == nullptr || peer_node == context.root_node || peer_node->op != GGML_OP_MUL_MAT_ID) { + return {}; + } + + size_t peer_index = 0; + if (!context.graph.index().node_index(peer_node, peer_index) || peer_index >= context.covered_nodes.size() || + context.covered_nodes[peer_index]) { + return {}; + } + + const std::vector & peer_consumers = context.graph.index().consumers(peer_output_id); + if (peer_consumers.size() != 1 || peer_consumers.front() != glu_node) { + return {}; + } + + const CommonMulMatIdMatch peer = common_match_mul_mat_id_any_format(context.graph, peer_node, context.plan); + if (!peer.matched() || peer.input->id != root.input->id || peer.route_ids->id != root.route_ids->id || + peer.input_size != root.input_size || peer.output_size != root.output_size || + peer.token_count != root.token_count || peer.route_count != root.route_count || + peer.input_route_count != root.input_route_count || peer.expert_count != root.expert_count || + !common_mul_mat_id_same_shape(*root.output, *peer.output)) { + return {}; + } + + const Value * output = common_mul_mat_id_graph_value(context.graph, glu_node->output); + if (output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || + !common_mul_mat_id_same_shape(*output, *root.output)) { + return {}; + } + + match.input = root.input; + match.route_ids = root.route_ids; + match.gate_weight = root_is_gate ? root.weight : peer.weight; + match.up_weight = root_is_gate ? peer.weight : root.weight; + match.gate_output = root_is_gate ? root.output : peer.output; + match.up_output = root_is_gate ? peer.output : root.output; + match.output = output; + match.gate_node = root_is_gate ? context.root_node : peer_node; + match.up_node = root_is_gate ? peer_node : context.root_node; + match.glu_node = glu_node; + match.routing_bundle = root.routing_bundle; + match.has_routing_bundle = root.has_routing_bundle; + match.input_size = root.input_size; + match.output_size = root.output_size; + match.token_count = root.token_count; + match.route_count = root.route_count; + match.input_route_count = root.input_route_count; + match.expert_count = root.expert_count; + match.gate_format = root_is_gate ? root.weight_format : peer.weight_format; + match.up_format = root_is_gate ? peer.weight_format : root.weight_format; + return match; +} + +static CommonMulMatIdMatch routing_match_for_swiglu(const MulMatIdSwiGLUMatch & match) { + CommonMulMatIdMatch routed; + routed.input = match.input; + routed.weight = match.gate_weight; + routed.output = match.gate_output; + routed.route_ids = match.route_ids; + routed.routing_bundle = match.routing_bundle; + routed.has_routing_bundle = match.has_routing_bundle; + routed.input_size = match.input_size; + routed.output_size = match.output_size; + routed.token_count = match.token_count; + routed.route_count = match.route_count; + routed.input_route_count = match.input_route_count; + routed.expert_count = match.expert_count; + routed.weight_format = match.gate_format; + return routed; +} + +static KernelCatalogRef mul_mat_id_swiglu_f32_kernel(const MulMatIdSwiGLUMatch & match) { + return match.token_count <= 5 ? kMulMatIdSkinnySwiGLUF32F32Kernel : kMulMatIdTiledSwiGLUF32F32Kernel; +} + +static bool match_mul_mat_id_swiglu_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatIdSwiGLUMatch match = match_mul_mat_id_swiglu(context); + if (!match.matched()) { + return false; + } + + const CommonMulMatIdMatch routed = routing_match_for_swiglu(match); + CommandPlanMoeRoutingBundle routing_bundle; + if (!common_mul_mat_id_ensure_moe_routing_bundle(context, routed, dispatch_match, routing_bundle)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(mul_mat_id_swiglu_f32_kernel(match)); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_mul_mat_id_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_swiglu.input_size", + common_mul_mat_id_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_swiglu.output_size", + common_mul_mat_id_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_swiglu.expert_count", + common_mul_mat_id_to_config_value(match.expert_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_swiglu.route_count", + common_mul_mat_id_to_config_value(match.route_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_swiglu.input_route_count", + common_mul_mat_id_to_config_value(match.input_route_count)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_id_swiglu.gate_weight_format", + common_mul_mat_id_to_config_value(common_mul_mat_format_config_value(match.gate_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_id_swiglu.up_weight_format", + common_mul_mat_id_to_config_value(common_mul_mat_format_config_value(match.up_format))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ routing_bundle.expert_table, 0, routing_bundle.expert_table_byte_count }); + dispatch.bindings.push_back({ routing_bundle.partition_table, 0, routing_bundle.partition_table_byte_count }); + dispatch.bindings.push_back({ match.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; +} + +} // namespace + +void register_gated_mul_mat_id_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.mul_mat_id_swiglu.f32_f32_wmma", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 200, + DispatchSource::Common, + match_mul_mat_id_swiglu_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.h new file mode 100644 index 000000000000..37cfb4dc9ada --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat-id.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_gated_mul_mat_id_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..c0016ba9f7d4 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp @@ -0,0 +1,1137 @@ +#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" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kMulMatF32F32WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatSwiGLUF32F32WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatTiledPairF32BinaryPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); +static constexpr KernelCatalogRef 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 = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output"); +static constexpr KernelCatalogRef kMulMatSwiGLUQ4Q8PrefillKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"); +static constexpr KernelCatalogRef kQuantizeF32SymmetricI4K32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_quantize_f32_symmetric_i4_k32"); +static constexpr KernelCatalogRef kMulMatSymmetricI4LowRowAdjacentDualWmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma"); +static constexpr KernelCatalogRef kMulMatSymmetricI4LowRowAdjacentDualDirectDotKernels[] = { + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c1"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c2"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c3"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c4"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c5"), +}; +static constexpr KernelCatalogRef kMulMatSwiGLUSymmetricI4WmmaQ8PlaneKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_swiglu_symmetric_i4_wmma_q8_plane"); +static constexpr KernelCatalogRef kMulMatQ5KQ8PlaneWmmaToken256Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q5_k_q8_plane_wmmai8_token256"); + +struct MulMatSwiGLUMatch { + const Value * input = nullptr; + const Value * gate_weight = nullptr; + const Value * up_weight = nullptr; + const Value * gate_output = nullptr; + const Value * up_output = nullptr; + const Value * output = nullptr; + const GraphNode * gate_node = nullptr; + const GraphNode * up_node = nullptr; + const GraphNode * glu_node = nullptr; + CommonMulMatWeightFormat gate_format = CommonMulMatWeightFormat::Q4K; + CommonMulMatWeightFormat up_format = CommonMulMatWeightFormat::Q4K; + BinaryKind op = BinaryKind::SwiGLU; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + + bool topology_matched() const { + return input != nullptr && gate_weight != nullptr && up_weight != nullptr && gate_output != nullptr && + up_output != nullptr && output != nullptr && gate_node != nullptr && up_node != nullptr && + glu_node != nullptr; + } + + bool matched() const { return topology_matched() && token_count >= 1; } + + bool decode_matched() const { return topology_matched() && token_count == 1; } +}; + +struct MulMatSwiGLUPostOpsMatch { + MulMatSwiGLUMatch root; + const Value * bias = nullptr; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const GraphNode * bias_add_node = nullptr; + const GraphNode * residual_add_node = nullptr; + std::vector add_nodes; + KernelCatalogRef kernel = {}; + bool has_bias = false; + bool has_residual = false; + + bool matched() const { + return root.matched() && residual_output != nullptr && kernel.id != kUncatalogedKernelId && + (has_bias || has_residual); + } +}; + +struct MulMatSwiGLUProjectionMatch { + MulMatSwiGLUMatch gate_up; + const GraphNode * projection_node = nullptr; + const Value * projection_weight = nullptr; + const Value * projection_output = nullptr; + int64_t projection_size = 0; + + bool matched() const { + return gate_up.matched() && projection_node != nullptr && projection_weight != nullptr && + projection_output != nullptr; + } +}; + +struct PackedMulMatGluMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * packed_output = nullptr; + const Value * output = nullptr; + const GraphNode * matmul_node = nullptr; + const GraphNode * glu_node = nullptr; + BinaryKind op = BinaryKind::SwiGLU; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + size_t gate_weight_offset = 0; + size_t up_weight_offset = 0; + size_t weight_half_bytes = 0; + CommonMulMatWeightFormat weight_format = CommonMulMatWeightFormat::Q4K; + + bool matched() const { + return input != nullptr && weight != nullptr && packed_output != nullptr && output != nullptr && + matmul_node != nullptr && glu_node != nullptr && token_count > 1; + } +}; + +static bool supported_symmetric_i4_pair(ggml_type gate_type, ggml_type up_type) { + return (gate_type == GGML_TYPE_Q4_K && up_type == GGML_TYPE_Q4_K) || + (gate_type == GGML_TYPE_Q5_K && up_type == GGML_TYPE_IQ4_XS) || + (gate_type == GGML_TYPE_IQ4_XS && up_type == GGML_TYPE_Q5_K); +} + +static bool supported_mixed_symmetric_i4_pair(ggml_type gate_type, ggml_type up_type) { + return (gate_type == GGML_TYPE_Q5_K && up_type == GGML_TYPE_IQ4_XS) || + (gate_type == GGML_TYPE_IQ4_XS && up_type == GGML_TYPE_Q5_K); +} + +static bool distinct_storage(const Graph & graph, const Value & lhs, const Value & rhs) { + return !graph.values().same_storage(lhs.id, rhs.id); +} + +static bool checked_mul_size(size_t lhs, size_t rhs, size_t & result) { + if (rhs != 0 && lhs > std::numeric_limits::max() / rhs) { + return false; + } + result = lhs * rhs; + return true; +} + +static bool is_bias_shape(const Value & value, int64_t output_size) { + if (value.kind != ValueKind::External || value.type != GGML_TYPE_F32 || !value.contiguous || + value.ne[0] != output_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] != 1) { + return false; + } + } + return true; +} + +static size_t symmetric_i4_weight_byte_count(int64_t input_size, int64_t output_size) { + return static_cast(output_size) * static_cast(input_size / 256) * size_t{ 144 }; +} + +static DispatchBinding symmetric_i4_weight_binding(const Value & weight, int64_t input_size, int64_t output_size) { + DispatchBinding binding; + binding.value = weight.id; + binding.length = symmetric_i4_weight_byte_count(input_size, output_size); + binding.layout = kSymmetricI4K32Row64Layout; + binding.source_type = weight.type; + binding.input_size = input_size; + binding.output_size = output_size; + binding.source_length = weight.byte_count; + return binding; +} + +static DispatchBinding symmetric_i5_weight_binding(const Value & weight, int64_t input_size, int64_t output_size) { + DispatchBinding binding; + binding.value = weight.id; + binding.length = weight.byte_count; + binding.layout = kQ5KSymmetricI5K32Layout; + binding.source_type = weight.type; + binding.input_size = input_size; + binding.output_size = output_size; + binding.source_length = weight.byte_count; + return binding; +} + +static MulMatSwiGLUMatch match_mul_mat_swiglu(const DispatchMatchContext & context) { + MulMatSwiGLUMatch match; + CommonMulMatMatch root = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + if (!root.matched()) { + root = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, true); + } + if (!root.matched() || !context.graph.has_index()) { + return match; + } + + const std::vector & root_consumers = context.graph.index().consumers(context.root_node->output); + if (root_consumers.size() != 1 || root_consumers.front() == nullptr) { + return {}; + } + + const GraphNode * glu_node = root_consumers.front(); + BinaryKind binary_op; + if (glu_node->inputs.size() != 2 || !common_fused_binary_kind_from_params(glu_node->params, binary_op)) { + return {}; + } + + size_t glu_index = 0; + if (!context.graph.index().node_index(glu_node, glu_index) || glu_index >= context.covered_nodes.size() || + context.covered_nodes[glu_index]) { + return {}; + } + + const bool root_is_gate = glu_node->inputs[0] == context.root_node->output; + const bool root_is_up = glu_node->inputs[1] == context.root_node->output; + if (!root_is_gate && !root_is_up) { + return {}; + } + + const ValueId peer_output_id = root_is_gate ? glu_node->inputs[1] : glu_node->inputs[0]; + const GraphNode * peer_node = context.graph.index().producer(peer_output_id); + if (peer_node == nullptr || peer_node == context.root_node || peer_node->op != GGML_OP_MUL_MAT) { + return {}; + } + + size_t peer_index = 0; + if (!context.graph.index().node_index(peer_node, peer_index) || peer_index >= context.covered_nodes.size() || + context.covered_nodes[peer_index]) { + return {}; + } + + const std::vector & peer_consumers = context.graph.index().consumers(peer_output_id); + if (peer_consumers.size() != 1 || peer_consumers.front() != glu_node) { + return {}; + } + + CommonMulMatMatch peer = common_match_mul_mat_any_format(context.graph, peer_node, kMulMatF32F32WmmaKernel, false); + if (!peer.matched()) { + peer = common_match_mul_mat_any_format(context.graph, peer_node, kMulMatF32F32WmmaKernel, true); + } + if (!peer.matched() || peer.input->id != root.input->id || peer.input_size != root.input_size || + peer.output_size != root.output_size || peer.token_count != root.token_count || + !common_same_shape(*root.output, *peer.output)) { + return {}; + } + + const Value * output = common_graph_value(context.graph, glu_node->output); + if (output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || + !common_same_shape(*output, *root.output)) { + return {}; + } + + match.input = root.input; + match.gate_weight = root_is_gate ? root.weight : peer.weight; + match.up_weight = root_is_gate ? peer.weight : root.weight; + match.gate_output = root_is_gate ? root.output : peer.output; + match.up_output = root_is_gate ? peer.output : root.output; + match.output = output; + match.gate_node = root_is_gate ? context.root_node : peer_node; + match.up_node = root_is_gate ? peer_node : context.root_node; + match.glu_node = glu_node; + match.gate_format = root_is_gate ? root.weight_format : peer.weight_format; + match.up_format = root_is_gate ? peer.weight_format : root.weight_format; + match.op = binary_op; + match.input_size = root.input_size; + match.output_size = root.output_size; + match.token_count = root.token_count; + return match; +} + +// Fuses a packed projection followed by one-input GLU. The packed weight is split into +// gate/up halves and lowered through the split-weight gated matmul kernel. +static PackedMulMatGluMatch match_packed_mul_mat_glu(const DispatchMatchContext & context) { + PackedMulMatGluMatch match; + const CommonMulMatMatch root = common_match_mul_mat_any_format(context.graph, context.root_node, + kMulMatTiledPairF32BinaryPublishF32Kernel, false); + if (!root.matched() || !context.graph.has_index() || root.weight->alias_source.value >= 0) { + return {}; + } + + const std::vector & root_consumers = context.graph.index().consumers(context.root_node->output); + if (root_consumers.size() != 1 || root_consumers.front() == nullptr || root_consumers.front()->op != GGML_OP_GLU) { + return {}; + } + + const GraphNode * glu_node = root_consumers.front(); + BinaryKind op; + if (glu_node->inputs.size() != 1 || glu_node->inputs[0] != context.root_node->output || + !common_fused_binary_kind_from_params(glu_node->params, op)) { + return {}; + } + + const GluParams * glu_params = op_params_as(glu_node->params); + if (glu_params == nullptr) { + return {}; + } + + size_t glu_index = 0; + if (!context.graph.index().node_index(glu_node, glu_index) || glu_index >= context.covered_nodes.size() || + context.covered_nodes[glu_index]) { + return {}; + } + + const Value * output = common_graph_value(context.graph, glu_node->output); + if (output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || root.output_size % 2 != 0 || + output->ne[0] != root.output_size / 2 || output->ne[1] != root.token_count || output->ne[2] != 1 || + output->ne[3] != 1 || output->alias_source.value >= 0) { + return {}; + } + + const size_t row_bytes = ggml_row_size(root.weight->type, root.input_size); + size_t half_bytes = 0; + if (row_bytes == 0 || !checked_mul_size(row_bytes, static_cast(output->ne[0]), half_bytes) || + half_bytes > root.weight->byte_count || half_bytes > std::numeric_limits::max() - half_bytes || + 2 * half_bytes > root.weight->byte_count) { + return {}; + } + + match.input = root.input; + match.weight = root.weight; + match.packed_output = root.output; + match.output = output; + match.matmul_node = context.root_node; + match.glu_node = glu_node; + match.op = op; + match.input_size = root.input_size; + match.output_size = output->ne[0]; + match.token_count = root.token_count; + match.gate_weight_offset = glu_params->swapped ? half_bytes : 0; + match.up_weight_offset = glu_params->swapped ? 0 : half_bytes; + match.weight_half_bytes = half_bytes; + match.weight_format = root.weight_format; + return match; +} + +static bool match_packed_mul_mat_glu_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const PackedMulMatGluMatch match = match_packed_mul_mat_glu(context); + if (!match.matched()) { + return false; + } + + 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); + 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", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.gate_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.weight_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.up_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.weight_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.weight->id, match.gate_weight_offset, match.weight_half_bytes }); + dispatch.bindings.push_back({ match.weight->id, match.up_weight_offset, match.weight_half_bytes }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.matmul_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; +} + +static bool match_mul_mat_swiglu_symmetric_i4_lowrow_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUMatch match = match_mul_mat_swiglu(context); + if (!match.topology_matched() || + !supported_mixed_symmetric_i4_pair(match.gate_weight->type, match.up_weight->type) || + match.gate_weight->alias_source.value >= 0 || match.up_weight->alias_source.value >= 0 || + match.token_count < 1 || match.token_count > 16 || match.input_size % 64 != 0 || match.output_size % 64 != 0 || + !distinct_storage(context.graph, *match.gate_weight, *match.up_weight) || + !distinct_storage(context.graph, *match.gate_output, *match.up_output)) { + return false; + } + + const CommonSymmetricI4ActivationLayout activation_layout = + common_symmetric_i4_activation_layout(match.input_size, match.token_count); + const CommandPlanAlternateValue * alternate = find_alternate_value(context.graph, context.plan, match.input->id, + GGML_TYPE_COUNT, activation_layout.total_bytes); + if (alternate != nullptr && alternate->name != kCommonSymmetricI4K32ActivationAlternateName) { + alternate = nullptr; + } + + const ValueId activation = alternate != nullptr ? alternate->alternate_value : context.next_plan_value; + if (alternate == nullptr) { + dispatch_match.transients.push_back( + { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeF32SymmetricI4K32Kernel); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.input_size", + common_to_config_value(match.input_size)); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.token_count", + common_to_config_value(match.token_count)); + quantize.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + quantize.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + quantize.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + quantize.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + dispatch_match.dispatches.push_back(std::move(quantize)); + + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.input->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, + kCommonSymmetricI4K32ActivationAlternateName }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + } + + const bool use_direct_dot = + match.token_count <= 5 && match.input_size % 256 == 0 && match.output_size >= 2 * match.input_size; + + Dispatch gate_up; + gate_up.kernel = make_kernel_specialization( + use_direct_dot ? kMulMatSymmetricI4LowRowAdjacentDualDirectDotKernels[match.token_count - 1] : + kMulMatSymmetricI4LowRowAdjacentDualWmmaKernel); + gate_up.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.input_size", + common_to_config_value(match.input_size)); + gate_up.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.output_size", + common_to_config_value(match.output_size)); + gate_up.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.token_count", + common_to_config_value(match.token_count)); + gate_up.kernel.compile_parameters.emplace( + "ggml.mul_mat.symmetric_i4.lowrow.row_group_size", + common_to_config_value( + static_cast(common_symmetric_shared4_row_group_size(match.input_size, match.output_size, 4)))); + gate_up.bindings.push_back( + match.gate_weight->type == GGML_TYPE_Q5_K && match.token_count == 1 ? + common_symmetric_i4_shared4_multistart_weight_binding(*match.gate_weight, match.input_size, + match.output_size) : + common_symmetric_i4_shared4_weight_binding(*match.gate_weight, match.input_size, match.output_size)); + gate_up.bindings.push_back( + match.up_weight->type == GGML_TYPE_Q5_K && match.token_count == 1 ? + common_symmetric_i4_shared4_multistart_weight_binding(*match.up_weight, match.input_size, + match.output_size) : + common_symmetric_i4_shared4_weight_binding(*match.up_weight, match.input_size, match.output_size)); + gate_up.bindings.push_back({ match.gate_output->id, 0, match.gate_output->byte_count }); + gate_up.bindings.push_back({ match.up_output->id, 0, match.up_output->byte_count }); + gate_up.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + gate_up.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + gate_up.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + + 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)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(gate_up)); + return dispatch_match.status.success(); +} + +static MulMatSwiGLUProjectionMatch match_mul_mat_swiglu_q5_projection(const DispatchMatchContext & context) { + MulMatSwiGLUProjectionMatch match; + match.gate_up = match_mul_mat_swiglu(context); + if (!match.gate_up.matched() || match.gate_up.op != BinaryKind::SwiGLU || + !supported_symmetric_i4_pair(match.gate_up.gate_weight->type, match.gate_up.up_weight->type) || + match.gate_up.gate_weight->alias_source.value >= 0 || match.gate_up.up_weight->alias_source.value >= 0 || + match.gate_up.token_count < 256 || match.gate_up.token_count % 256 != 0 || + match.gate_up.output_size % 256 != 0 || + !distinct_storage(context.graph, *match.gate_up.gate_weight, *match.gate_up.up_weight)) { + return {}; + } + + const GraphNode * projection_node = + common_find_only_consumer_with_op(context.graph, match.gate_up.output->id, GGML_OP_MUL_MAT); + if (projection_node == nullptr || projection_node->inputs.size() != 2 || + projection_node->inputs[1] != match.gate_up.output->id) { + return {}; + } + + size_t projection_index = 0; + if (!context.graph.index().node_index(projection_node, projection_index) || + projection_index >= context.covered_nodes.size() || context.covered_nodes[projection_index]) { + return {}; + } + + const CommonMulMatMatch projection = + common_match_mul_mat_any_format(context.graph, projection_node, kMulMatQ5KQ8PlaneWmmaToken256Kernel, false); + if (!projection.matched() || projection.weight->type != GGML_TYPE_Q5_K || + projection.weight->alias_source.value >= 0 || projection.input->id != match.gate_up.output->id || + projection.input_size != match.gate_up.output_size || projection.token_count != match.gate_up.token_count || + projection.output_size % 64 != 0 || !distinct_storage(context.graph, *projection.weight, *projection.output)) { + return {}; + } + + match.projection_node = projection_node; + match.projection_weight = projection.weight; + match.projection_output = projection.output; + match.projection_size = projection.output_size; + return match; +} + +static bool match_mul_mat_swiglu_q5_projection_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUProjectionMatch match = match_mul_mat_swiglu_q5_projection(context); + if (!match.matched()) { + return false; + } + + const size_t input_elements = + static_cast(match.gate_up.token_count) * static_cast(match.gate_up.input_size); + const size_t i4_payload_bytes = input_elements / 2; + const size_t i4_metadata_bytes = input_elements / 8; + const size_t q8_output_bytes = + static_cast(match.gate_up.token_count) * ggml_row_size(GGML_TYPE_Q8_1, match.gate_up.output_size); + const ValueId i4_payload = context.next_plan_value; + const ValueId i4_scales(context.next_plan_value.value + 1); + const ValueId i4_sums(context.next_plan_value.value + 2); + const ValueId q8_output(context.next_plan_value.value + 3); + + dispatch_match.transients.push_back( + { i4_payload, "common.mul_mat_swiglu.symmetric_i4.payload", i4_payload_bytes, 256 }); + dispatch_match.transients.push_back( + { i4_scales, "common.mul_mat_swiglu.symmetric_i4.scales", i4_metadata_bytes, 256 }); + dispatch_match.transients.push_back({ i4_sums, "common.mul_mat_swiglu.symmetric_i4.sums", i4_metadata_bytes, 256 }); + dispatch_match.transients.push_back({ q8_output, "common.mul_mat_swiglu.q8_plane", q8_output_bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeF32SymmetricI4K32Kernel); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.input_size", + common_to_config_value(match.gate_up.input_size)); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.token_count", + common_to_config_value(match.gate_up.token_count)); + quantize.bindings.push_back({ match.gate_up.input->id, 0, match.gate_up.input->byte_count }); + quantize.bindings.push_back({ i4_payload, 0, i4_payload_bytes }); + quantize.bindings.push_back({ i4_scales, 0, i4_metadata_bytes }); + quantize.bindings.push_back({ i4_sums, 0, i4_metadata_bytes }); + + Dispatch gate_up; + gate_up.kernel = make_kernel_specialization(kMulMatSwiGLUSymmetricI4WmmaQ8PlaneKernel); + gate_up.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.symmetric_i4.input_size", + common_to_config_value(match.gate_up.input_size)); + gate_up.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.symmetric_i4.output_size", + common_to_config_value(match.gate_up.output_size)); + gate_up.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.symmetric_i4.token_count", + common_to_config_value(match.gate_up.token_count)); + gate_up.bindings.push_back( + symmetric_i4_weight_binding(*match.gate_up.gate_weight, match.gate_up.input_size, match.gate_up.output_size)); + gate_up.bindings.push_back( + symmetric_i4_weight_binding(*match.gate_up.up_weight, match.gate_up.input_size, match.gate_up.output_size)); + gate_up.bindings.push_back({ match.gate_up.input->id, 0, match.gate_up.input->byte_count }); + gate_up.bindings.push_back({ match.gate_up.output->id, 0, match.gate_up.output->byte_count }); + gate_up.bindings.push_back({ i4_payload, 0, i4_payload_bytes }); + gate_up.bindings.push_back({ i4_scales, 0, i4_metadata_bytes }); + gate_up.bindings.push_back({ i4_sums, 0, i4_metadata_bytes }); + gate_up.bindings.push_back({ q8_output, 0, q8_output_bytes }); + + Dispatch projection; + projection.kernel = make_kernel_specialization(kMulMatQ5KQ8PlaneWmmaToken256Kernel); + projection.kernel.integer_parameters.emplace("token_count", match.gate_up.token_count); + projection.kernel.compile_parameters.emplace("ggml.mul_mat_q5_k_q8_plane.input_size", + common_to_config_value(match.gate_up.output_size)); + projection.kernel.compile_parameters.emplace("ggml.mul_mat_q5_k_q8_plane.output_size", + common_to_config_value(match.projection_size)); + projection.kernel.compile_parameters.emplace("ggml.mul_mat_q5_k_q8_plane.token_capacity", + common_to_config_value(match.gate_up.token_count)); + projection.bindings.push_back({ q8_output, 0, q8_output_bytes }); + projection.bindings.push_back( + symmetric_i5_weight_binding(*match.projection_weight, match.gate_up.output_size, match.projection_size)); + projection.bindings.push_back({ match.projection_output->id, 0, match.projection_output->byte_count }); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_up.gate_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_up.up_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_up.glu_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.projection_node, + dispatch_match.covered_nodes)) { + return false; + } + + dispatch_match.dispatches.push_back(std::move(quantize)); + dispatch_match.dispatches.push_back(std::move(gate_up)); + dispatch_match.dispatches.push_back(std::move(projection)); + return true; +} + +static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUMatch match = match_mul_mat_swiglu(context); + if (!match.matched() || match.gate_format != CommonMulMatWeightFormat::Q4K || + match.up_format != CommonMulMatWeightFormat::Q4K || match.token_count < 256 || match.token_count > 2048 || + match.token_count % 256 != 0 || match.input_size % 256 != 0 || match.output_size % 64 != 0 || + match.gate_weight->alias_source.value >= 0 || match.up_weight->alias_source.value >= 0 || + !distinct_storage(context.graph, *match.gate_weight, *match.up_weight) || + !distinct_storage(context.graph, *match.gate_output, *match.up_output)) { + return false; + } + + // Keep in sync with ggml_swiglu_use_packed_f16 in the shared operation. + const bool packed_input = + match.token_count % 512 == 0 && (match.token_count / 512) * (match.output_size / 64) >= 64; + DispatchBinding activation; + const bool prepared = packed_input ? + common_prepare_k16_major_f16_input(context, *match.input, match.input_size, + match.token_count, dispatch_match, activation) : + common_prepare_f16_input(context, *match.input, match.input_size, match.token_count, + dispatch_match, activation); + if (!prepared) { + return false; + } + + 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); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.f16_output_layout", packed_output ? "1" : "0"); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.bindings.push_back(activation); + for (const Value * weight : { match.gate_weight, match.up_weight }) { + dispatch.bindings.push_back({ weight->id, 0, weight->byte_count, kQ4KPackedK256Row64Layout, weight->type, + match.input_size, match.output_size, weight->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + const char * f16_name = packed_output ? "common.mul_mat_swiglu.k16_major_f16" : "common.mul_mat_swiglu.f16"; + 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)) { + return false; + } + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool supports_swiglu_q8_output_shape(const MulMatSwiGLUMatch & match) { + return match.op == BinaryKind::SwiGLU && match.token_count >= 1 && match.token_count <= 5 && + common_is_supported_dense_input_size(CommonMulMatWeightFormat::Q4K, match.input_size) && + common_is_supported_dense_output_size(match.output_size) && match.output_size % 128 == 0; +} + +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) { + if (!match.matched()) { + return false; + } + + const bool pack_q4 = match.token_count <= 5 && match.output_size % 64 == 0 && + match.gate_weight->alias_source.value < 0 && match.up_weight->alias_source.value < 0 && + match.gate_format == CommonMulMatWeightFormat::Q4K && + match.up_format == CommonMulMatWeightFormat::Q4K; + const bool use_q8 = pack_q4; + const bool use_direct_dot = match.token_count <= 5 && match.input_size % 256 == 0; + return !use_q8 && !use_direct_dot && match.token_count >= 2; +} + +static 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 {}; + } + 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) { + MulMatSwiGLUPostOpsMatch match; + match.root = match_mul_mat_swiglu(context); + if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match.root) || !context.graph.has_index()) { + return {}; + } + + const Value * current = match.root.output; + for (int add_index = 0; add_index < 2; ++add_index) { + const GraphNode * add_node = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_ADD); + if (add_node == nullptr || !common_binary_node_is_add(*add_node)) { + break; + } + + const bool current_is_lhs = add_node->inputs[0] == current->id; + const bool current_is_rhs = add_node->inputs[1] == current->id; + if (!current_is_lhs && !current_is_rhs) { + return {}; + } + + const ValueId other_id = current_is_lhs ? add_node->inputs[1] : add_node->inputs[0]; + const Value * other = common_graph_value(context.graph, other_id); + const Value * output = common_graph_value(context.graph, add_node->output); + if (other == nullptr || output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || + !common_same_shape(*match.root.output, *output)) { + return {}; + } + + if (!match.has_bias && is_bias_shape(*other, match.root.output_size)) { + match.bias = other; + match.bias_add_node = add_node; + match.add_nodes.push_back(add_node); + match.has_bias = true; + current = output; + continue; + } + + if (!match.has_residual && other->type == GGML_TYPE_F32 && other->contiguous && + common_same_shape(*match.root.output, *other)) { + match.residual_input = other; + match.residual_output = output; + match.residual_add_node = add_node; + match.add_nodes.push_back(add_node); + match.has_residual = true; + current = output; + continue; + } + + break; + } + + if (!match.has_bias && !match.has_residual) { + return {}; + } + + match.residual_output = current; + match.kernel = kMulMatTiledPairF32BinaryBiasAddPublishF32Kernel; + return match; +} + +static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, + const MulMatSwiGLUPostOpsMatch * postops, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUMatch & match = postops != nullptr ? postops->root : match_mul_mat_swiglu(context); + if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { + return false; + } + + 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); + 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", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.gate_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.gate_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.up_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.up_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + if (postops != nullptr) { + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.epilogue", + common_to_config_value((postops->has_bias ? int64_t{ 1 } : int64_t{ 0 }) | + (postops->has_residual ? int64_t{ 2 } : int64_t{ 0 }))); + } + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); + dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + const Value * output = postops != nullptr ? postops->residual_output : match.output; + if (postops != nullptr) { + const Value * placeholder = postops->has_bias ? postops->bias : postops->residual_input; + dispatch.bindings.push_back(postops->has_bias ? + DispatchBinding{ postops->bias->id, 0, postops->bias->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back( + postops->has_residual ? + DispatchBinding{ postops->residual_input->id, 0, postops->residual_input->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + } + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + if (publish_k16) { + 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(publication.binding()); + } + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.up_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.glu_node, + dispatch_match.covered_nodes)) { + return false; + } + if (postops != nullptr) { + for (const GraphNode * add_node : postops->add_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, add_node, + dispatch_match.covered_nodes)) { + return false; + } + } + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_tiled_pair_mul_mat_swiglu_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatSwiGLUPostOpsMatch match = match_mul_mat_swiglu_postops(context); + if (!match.matched()) { + return false; + } + return build_tiled_pair_mul_mat_swiglu_dispatch(context, &match, dispatch_match); +} + +static bool match_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + return build_tiled_pair_mul_mat_swiglu_dispatch(context, nullptr, dispatch_match); +} + +static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatSwiGLUMatch match = match_mul_mat_swiglu(context); + if (!match.matched()) { + return false; + } + if (common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { + return false; + } + + const bool 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 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, + activation, CommonQ8ActivationPolicy::AllowStandaloneQuantize)) { + return false; + } + + Dispatch dispatch; + const bool use_direct_dot = match.token_count <= 5 && match.input_size % 256 == 0; + const bool use_wmma = !use_q8 && !use_direct_dot; + if (use_wmma && match.token_count < 2) { + return false; + } + dispatch.kernel = make_kernel_specialization( + publish_q8 ? kMulMatSwiGLUQ4Q8OutputKernel : + use_q8 ? kMulMatSwiGLUQ4Q8LowTokenDotKernel : + (use_direct_dot ? kMulMatSwiGLUF32F32LowTokenDotKernel : kMulMatSwiGLUF32F32WmmaKernel)); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.gate_weight_format", + common_to_config_value(common_mul_mat_format_config_value( + pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : match.gate_format))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.up_weight_format", + common_to_config_value(common_mul_mat_format_config_value( + pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : match.up_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.bindings.push_back(activation); + if (pack_q4) { + for (const Value * weight : { match.gate_weight, match.up_weight }) { + dispatch.bindings.push_back({ weight->id, 0, weight->byte_count, kQ4KPackedK256Row64Layout, weight->type, + match.input_size, match.output_size, weight->byte_count }); + } + } else { + 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 (publish_q8) { + constexpr const char * name = "common.mul_mat_swiglu.q8_1_x4"; + 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(publication.binding()); + } + + 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; +} + +static bool match_decode_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatSwiGLUMatch match = match_mul_mat_swiglu(context); + if (!match.decode_matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUF32F32DecodeWave64Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.gate_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.gate_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.up_weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.up_format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); + dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + 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; +} + +} // namespace + +void register_gated_mul_mat_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.mul_mat_swiglu.q4_k_q8_1_x4_prefill", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 310, + DispatchSource::Common, + match_mul_mat_swiglu_q4_q8_prefill_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu.symmetric_i4_lowrow_adjacent_dual", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 310, + DispatchSource::Common, + match_mul_mat_swiglu_symmetric_i4_lowrow_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu_q5_projection.symmetric_i4_q8_plane", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_mul_mat_swiglu_q5_projection_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 292, + DispatchSource::Common, + match_tiled_pair_mul_mat_swiglu_postops_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu.tiled_pair_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 291, + DispatchSource::Common, + match_tiled_pair_mul_mat_swiglu_dispatch, + }); + registry.add({ + "common.mul_mat_swiglu.legacy_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 290, + DispatchSource::Common, + match_mul_mat_swiglu_dispatch, + }); + registry.add({ + "common.packed_mul_mat_glu.tiled_pair_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 285, + DispatchSource::Common, + match_packed_mul_mat_glu_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.h new file mode 100644 index 000000000000..07e1477be1f8 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_gated_mul_mat_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp new file mode 100644 index 000000000000..d7a49803fa75 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.cpp @@ -0,0 +1,425 @@ +#include "dispatch-gather-add.h" + +#include "dispatch-binary-common.h" +#include "dispatch-mul-mat-common.h" +#include "ggml.h" +#include "graph/op-params.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kGatherAddF32Kernel = GGML_HRX_KERNEL_REF("hrx", "ggml_gather_add_f32"); +static constexpr KernelCatalogRef kGatherAddRmsNormBinaryF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_gather_add_rmsnorm_binary_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_2d_f32(const Value & value) { + return value.type == GGML_TYPE_F32 && value.contiguous && value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && + value.ne[3] == 1; +} + +static bool is_row_id_tensor(const Value & value, int64_t output_token_count) { + return value.type == GGML_TYPE_I32 && value.contiguous && value.element_count == output_token_count && + output_token_count > 0; +} + +static bool is_supported_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +static bool is_supported_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +static bool value_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +struct GatherAddMatch { + const GraphNode * first_get_rows = nullptr; + const GraphNode * second_get_rows = nullptr; + const GraphNode * add_node = nullptr; + const Value * first_source = nullptr; + const Value * second_source = nullptr; + const Value * row_ids = nullptr; + const Value * output = nullptr; + size_t first_get_rows_index = 0; + size_t second_get_rows_index = 0; + size_t add_node_index = 0; + int64_t source_token_count = 0; + int64_t output_token_count = 0; + int64_t hidden_size = 0; + BinaryKind op = BinaryKind::Add; + + bool matched() const { + return first_get_rows != nullptr && second_get_rows != nullptr && add_node != nullptr && + first_source != nullptr && second_source != nullptr && row_ids != nullptr && output != nullptr; + } +}; + +struct GatherAddRmsNormBinaryMatch { + GatherAddMatch gather_add; + const GraphNode *rms_node = nullptr; + const GraphNode *binary_node = nullptr; + const Value * weight = nullptr; + const Value * normalized_output = nullptr; + size_t rms_node_index = 0; + size_t binary_node_index = 0; + float epsilon = 0.0f; + + bool matched() const { + return gather_add.matched() && rms_node != nullptr && binary_node != nullptr && weight != nullptr && + normalized_output != nullptr; + } +}; + +static const GraphNode * find_single_binary_consumer(const Graph & graph, const GraphNode & get_rows) { + const std::vector & consumers = graph.index().consumers(get_rows.output); + if (consumers.size() != 1) { + return nullptr; + } + const GraphNode * binary = consumers.front(); + return common_match_binary_consumer(binary, get_rows.output, kCommonArithmeticBinaryKinds).matched() ? binary : + nullptr; +} + +static const GraphNode * peer_get_rows_input(const Graph & graph, + const GraphNode & add, + const GraphNode & root_get_rows) { + ValueId peer_output; + if (add.inputs[0] == root_get_rows.output) { + peer_output = add.inputs[1]; + } else if (add.inputs[1] == root_get_rows.output) { + peer_output = add.inputs[0]; + } else { + return nullptr; + } + + const GraphNode * peer = graph.index().producer(peer_output); + return peer != nullptr && peer->op == GGML_OP_GET_ROWS && peer->inputs.size() == 2 ? peer : nullptr; +} + +static GatherAddMatch match_gather_add_f32(const Graph & graph, const GraphNode * node, size_t node_index) { + GatherAddMatch match; + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2 || !graph.has_index()) { + return match; + } + + const GraphNode * add_node = find_single_binary_consumer(graph, *node); + const GraphNode * peer = add_node != nullptr ? peer_get_rows_input(graph, *add_node, *node) : nullptr; + if (add_node == nullptr || peer == nullptr) { + return {}; + } + + size_t peer_index = 0; + size_t add_index = 0; + if (!graph.index().node_index(peer, peer_index) || !graph.index().node_index(add_node, add_index)) { + return {}; + } + + const GraphNode * first_get_rows = add_node->inputs[0] == node->output ? node : peer; + const GraphNode * second_get_rows = first_get_rows == node ? peer : node; + const size_t first_get_rows_index = first_get_rows == node ? node_index : peer_index; + const size_t second_get_rows_index = first_get_rows == node ? peer_index : node_index; + const Value * first_source = graph_value(graph, first_get_rows->inputs[0]); + const Value * first_ids = graph_value(graph, first_get_rows->inputs[1]); + const Value * first_output = graph_value(graph, first_get_rows->output); + const Value * second_source = graph_value(graph, second_get_rows->inputs[0]); + const Value * second_ids = graph_value(graph, second_get_rows->inputs[1]); + const Value * second_output = graph_value(graph, second_get_rows->output); + const Value * output = graph_value(graph, add_node->output); + if (first_source == nullptr || first_ids == nullptr || first_output == nullptr || second_source == nullptr || + second_ids == nullptr || second_output == nullptr || output == nullptr) { + return {}; + } + if (first_get_rows->inputs[1] != second_get_rows->inputs[1]) { + return {}; + } + if (!is_2d_f32(*first_source) || !is_2d_f32(*second_source) || !is_2d_f32(*first_output) || + !is_2d_f32(*second_output) || !is_2d_f32(*output)) { + return {}; + } + if (!same_shape(*first_source, *second_source) || !same_shape(*first_output, *second_output) || + !same_shape(*first_output, *output)) { + return {}; + } + + const int64_t hidden_size = first_source->ne[0]; + const int64_t source_token_count = first_source->ne[1]; + const int64_t output_token_count = first_output->ne[1]; + if (first_output->ne[0] != hidden_size || !is_row_id_tensor(*first_ids, output_token_count) || + !is_row_id_tensor(*second_ids, output_token_count)) { + return {}; + } + if (!is_supported_hidden_size(hidden_size) || !is_supported_token_count(source_token_count) || + !is_supported_token_count(output_token_count)) { + return {}; + } + + const CommonBinaryMatch binary = + common_match_binary_consumer(add_node, first_get_rows->output, kCommonArithmeticBinaryKinds); + if (!binary.matched() || !binary.producer_is_lhs) { + return {}; + } + + match.first_get_rows = first_get_rows; + match.second_get_rows = second_get_rows; + match.add_node = add_node; + match.first_source = first_source; + match.second_source = second_source; + match.row_ids = first_ids; + match.output = output; + match.first_get_rows_index = first_get_rows_index; + match.second_get_rows_index = second_get_rows_index; + match.add_node_index = add_index; + match.source_token_count = source_token_count; + match.output_token_count = output_token_count; + match.hidden_size = hidden_size; + match.op = binary.kind; + return match; +} + +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) { + for (const GraphNode * consumer : graph.index().consumers(value.id)) { + if (is_packed_q8_consumer(graph, consumer, value)) { + return true; + } + } + return false; +} + +template static bool pairwise_distinct_storage_roots(const std::array & values) { + for (size_t lhs = 0; lhs < values.size(); ++lhs) { + for (size_t rhs = lhs + 1; rhs < values.size(); ++rhs) { + if (values[lhs]->storage_root == values[rhs]->storage_root) { + return false; + } + } + } + return true; +} + +static GatherAddRmsNormBinaryMatch match_gather_add_rmsnorm_binary_f32(const Graph & graph, + const GraphNode * node, + size_t node_index) { + GatherAddRmsNormBinaryMatch match; + match.gather_add = match_gather_add_f32(graph, node, node_index); + if (!match.gather_add.matched()) { + return {}; + } + + const GraphNode * rms_node = nullptr; + for (const GraphNode * consumer : graph.index().consumers(match.gather_add.output->id)) { + if (consumer != nullptr && consumer->op == GGML_OP_RMS_NORM && consumer->inputs.size() == 1) { + if (rms_node != nullptr) { + return {}; + } + rms_node = consumer; + } + } + size_t rms_node_index = 0; + if (rms_node == nullptr || !graph.index().node_index(rms_node, rms_node_index)) { + return {}; + } + const RmsNormParams * rms_params = op_params_as(rms_node->params); + const Value * rms_output = graph_value(graph, rms_node->output); + if (rms_params == nullptr || !std::isfinite(rms_params->eps) || rms_params->eps <= 0.0f || rms_output == nullptr || + rms_output->type != GGML_TYPE_F32 || !rms_output->contiguous || + !same_shape(*match.gather_add.output, *rms_output)) { + return {}; + } + const std::vector & rms_consumers = graph.index().consumers(rms_output->id); + if (rms_consumers.size() != 1 || rms_consumers.front() == nullptr || rms_consumers.front()->op != GGML_OP_MUL || + rms_consumers.front()->inputs.size() != 2) { + return {}; + } + const GraphNode * binary_node = rms_consumers.front(); + size_t binary_node_index = 0; + if (!graph.index().node_index(binary_node, binary_node_index)) { + return {}; + } + const ValueId weight_id = binary_node->inputs[0] == rms_output->id ? binary_node->inputs[1] : + binary_node->inputs[1] == rms_output->id ? binary_node->inputs[0] : + ValueId{}; + const Value * weight = graph_value(graph, weight_id); + const Value * normalized_output = graph_value(graph, binary_node->output); + if (weight == nullptr || normalized_output == nullptr || weight->type != GGML_TYPE_F32 || + normalized_output->type != GGML_TYPE_F32 || !weight->contiguous || !normalized_output->contiguous || + weight->ne[0] != match.gather_add.hidden_size || weight->ne[1] != 1 || weight->ne[2] != 1 || + weight->ne[3] != 1 || !same_shape(*match.gather_add.output, *normalized_output) || + has_packed_q8_consumer(graph, *normalized_output) || + common_has_symmetric_i4_lowrow_consumer(graph, *normalized_output) || + !pairwise_distinct_storage_roots(std::array{ + match.gather_add.first_source, match.gather_add.second_source, match.gather_add.row_ids, + match.gather_add.output, weight, normalized_output })) { + return {}; + } + + match.rms_node = rms_node; + match.binary_node = binary_node; + match.weight = weight; + match.normalized_output = normalized_output; + match.rms_node_index = rms_node_index; + match.binary_node_index = binary_node_index; + match.epsilon = rms_params->eps; + return match; +} + +static std::string to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +} // namespace + +static bool match_gather_add_rmsnorm_binary_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GatherAddRmsNormBinaryMatch fused = + match_gather_add_rmsnorm_binary_f32(context.graph, context.root_node, context.root_index); + if (!fused.matched() || fused.gather_add.first_get_rows_index >= context.covered_nodes.size() || + fused.gather_add.second_get_rows_index >= context.covered_nodes.size() || + fused.gather_add.add_node_index >= context.covered_nodes.size() || + fused.rms_node_index >= context.covered_nodes.size() || + fused.binary_node_index >= context.covered_nodes.size() || + context.covered_nodes[fused.gather_add.first_get_rows_index] || + context.covered_nodes[fused.gather_add.second_get_rows_index] || + context.covered_nodes[fused.gather_add.add_node_index] || context.covered_nodes[fused.rms_node_index] || + context.covered_nodes[fused.binary_node_index]) { + return false; + } + if (!value_is_available(context.graph, fused.gather_add.first_source->id, context.covered_nodes) || + !value_is_available(context.graph, fused.gather_add.second_source->id, context.covered_nodes) || + !value_is_available(context.graph, fused.gather_add.row_ids->id, context.covered_nodes) || + !value_is_available(context.graph, fused.weight->id, context.covered_nodes)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGatherAddRmsNormBinaryF32Kernel); + dispatch.kernel.integer_parameters.emplace("source_token_count", fused.gather_add.source_token_count); + dispatch.kernel.integer_parameters.emplace("output_token_count", fused.gather_add.output_token_count); + dispatch.kernel.compile_parameters.emplace("ggml.gather_add_rmsnorm_binary_f32.hidden_size", + std::to_string(fused.gather_add.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", + to_config_value(fused.epsilon)); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.gather_add_rmsnorm_binary_f32.gather_op", nullptr, + fused.gather_add.op, true); + dispatch.bindings.push_back({ fused.gather_add.first_source->id, 0, fused.gather_add.first_source->byte_count }); + dispatch.bindings.push_back({ fused.gather_add.second_source->id, 0, fused.gather_add.second_source->byte_count }); + dispatch.bindings.push_back({ fused.gather_add.row_ids->id, 0, fused.gather_add.row_ids->byte_count }); + dispatch.bindings.push_back({ fused.gather_add.output->id, 0, fused.gather_add.output->byte_count }); + dispatch.bindings.push_back({ fused.weight->id, 0, fused.weight->byte_count }); + dispatch.bindings.push_back({ fused.normalized_output->id, 0, fused.normalized_output->byte_count }); + + match.covered_nodes.push_back(fused.gather_add.first_get_rows_index); + match.covered_nodes.push_back(fused.gather_add.second_get_rows_index); + match.covered_nodes.push_back(fused.gather_add.add_node_index); + 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)); + return true; +} + +static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GatherAddMatch gather_add = match_gather_add_f32(context.graph, context.root_node, context.root_index); + if (!gather_add.matched() || gather_add.first_get_rows_index >= context.covered_nodes.size() || + gather_add.second_get_rows_index >= context.covered_nodes.size() || + gather_add.add_node_index >= context.covered_nodes.size() || + context.covered_nodes[gather_add.first_get_rows_index] || + context.covered_nodes[gather_add.second_get_rows_index] || context.covered_nodes[gather_add.add_node_index]) { + return false; + } + if (!value_is_available(context.graph, gather_add.first_source->id, context.covered_nodes) || + !value_is_available(context.graph, gather_add.second_source->id, context.covered_nodes) || + !value_is_available(context.graph, gather_add.row_ids->id, context.covered_nodes)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGatherAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("source_token_count", gather_add.source_token_count); + dispatch.kernel.integer_parameters.emplace("output_token_count", gather_add.output_token_count); + dispatch.kernel.integer_parameters.emplace("hidden_size", gather_add.hidden_size); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.gather_add_f32.binary_op", nullptr, gather_add.op, + true); + dispatch.bindings.push_back({ gather_add.first_source->id, 0, gather_add.first_source->byte_count }); + dispatch.bindings.push_back({ gather_add.second_source->id, 0, gather_add.second_source->byte_count }); + dispatch.bindings.push_back({ gather_add.row_ids->id, 0, gather_add.row_ids->byte_count }); + dispatch.bindings.push_back({ gather_add.output->id, 0, gather_add.output->byte_count }); + + match.covered_nodes.push_back(gather_add.first_get_rows_index); + match.covered_nodes.push_back(gather_add.second_get_rows_index); + match.covered_nodes.push_back(gather_add.add_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_gather_add_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.gather_add_rmsnorm_binary_f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 1100, + DispatchSource::Common, + match_gather_add_rmsnorm_binary_f32_dispatch, + }); + registry.add({ + "common.gather_add_f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Common, + match_gather_add_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.h similarity index 80% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.h rename to ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.h index b8d43035350f..9f8d6e94a4d0 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.h +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gather-add.h @@ -1,6 +1,6 @@ #pragma once -#include "dispatch-registry.h" +#include "../dispatch-registry.h" namespace ggml::hrx { 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 new file mode 100644 index 000000000000..022d1f05279e --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.cpp @@ -0,0 +1,509 @@ +#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 +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kGetRowsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_f32"); +static constexpr KernelCatalogRef kGetRowsF32NextKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_f32_next"); +static constexpr KernelCatalogRef kGetRowsScaleF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_get_rows_scale_f32"); +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 kMaxGetRowsRowCount = 262208; + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_1d_or_2d_column(const Value & value) { + return value.ne[0] > 0 && value.ne[1] == 1 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_supported_hidden_size(ggml_type type, int64_t hidden_size) { + if (hidden_size < 4 || hidden_size > kMaximumHiddenElements || hidden_size % 4 != 0) { + return false; + } + + const int64_t block_size = ggml_blck_size(type); + return block_size > 0 && hidden_size % block_size == 0; +} + +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; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static std::string to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +static bool get_rows_format_for_type(ggml_type type, int64_t & format) { + if (type == GGML_TYPE_Q1_0) { + format = kQ1_0GetRowsFormat; + return true; + } + + CommonMulMatWeightFormat common_format; + if (!common_mul_mat_format_for_type(type, common_format)) { + return false; + } + + format = common_mul_mat_format_config_value(common_format); + return true; +} + +struct GetRowsMatch { + const Value * ids = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + int64_t weight_format_value = -1; + int64_t token_count = 0; + int64_t row_count = 0; + int64_t hidden_size = 0; + + bool matched() const { + return ids != nullptr && weight != nullptr && output != nullptr && weight_format_value >= 0; + } +}; + +struct GetRowsRmsNormMatch { + GetRowsMatch rows; + const GraphNode * rms_node = nullptr; + const GraphNode * binary_node = nullptr; + const Value * norm_weight = nullptr; + const Value * normalized_output = nullptr; + size_t rms_node_index = 0; + size_t binary_node_index = 0; + float epsilon = 0.0f; + + bool matched() const { + return rows.matched() && rms_node != nullptr && binary_node != nullptr && norm_weight != nullptr && + normalized_output != nullptr; + } +}; + +struct GetRowsScaleMatch { + GetRowsMatch rows; + const GraphNode * scale_node = nullptr; + const Value * output = nullptr; + size_t scale_index = 0; + float scale = 0.0f; + float bias = 0.0f; + + bool matched() const { + return rows.matched() && scale_node != nullptr && output != nullptr; + } +}; + +static void add_common_compile_parameters(Dispatch & dispatch, const GetRowsMatch & match) { + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.token_capacity", to_config_value(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)); +} + +static void add_common_integer_parameters(Dispatch & dispatch, const GetRowsMatch & match) { + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.integer_parameters.emplace("row_count", match.row_count); + dispatch.kernel.integer_parameters.emplace("hidden_size", match.hidden_size); +} + +static void add_primary_bindings(Dispatch & dispatch, const GetRowsMatch & match) { + dispatch.bindings.push_back({ match.ids->id, 0, match.ids->byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); +} + +static GetRowsMatch match_get_rows_f32(const Graph & graph, const GraphNode * node) { + GetRowsMatch match; + if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * ids = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (weight == nullptr || ids == nullptr || output == nullptr || weight->type == GGML_TYPE_IQ3_S || + weight->type == GGML_TYPE_IQ4_NL || ids->type != GGML_TYPE_I32 || output->type != GGML_TYPE_F32 || + !weight->contiguous || !ids->contiguous || !output->contiguous || !is_2d(*weight) || + !is_1d_or_2d_column(*ids) || !is_2d(*output)) { + return {}; + } + + int64_t weight_format_value = 0; + if (!get_rows_format_for_type(weight->type, weight_format_value)) { + return {}; + } + + const int64_t hidden_size = weight->ne[0]; + const int64_t row_count = weight->ne[1]; + const int64_t token_count = ids->ne[0]; + if (output->ne[0] != hidden_size || output->ne[1] != token_count || + !is_supported_hidden_size(weight->type, hidden_size) || !is_supported_row_count(row_count) || + !is_supported_token_count(token_count)) { + return {}; + } + + match.ids = ids; + match.weight = weight; + match.output = output; + match.weight_format_value = weight_format_value; + match.token_count = token_count; + match.row_count = row_count; + match.hidden_size = hidden_size; + return match; +} + +static GetRowsScaleMatch match_get_rows_scale_f32(const DispatchMatchContext & context) { + GetRowsScaleMatch match; + match.rows = match_get_rows_f32(context.graph, context.root_node); + if (!match.rows.matched() || !context.graph.has_index()) { + return {}; + } + + const bool supported_weight = match.rows.weight->type == GGML_TYPE_F32 || + (match.rows.weight->type == GGML_TYPE_Q5_1 && match.rows.token_count > 1); + if (!supported_weight) { + return {}; + } + + const std::vector & consumers = context.graph.index().consumers(match.rows.output->id); + if (consumers.size() != 1 || consumers.front() == nullptr || consumers.front()->op != GGML_OP_SCALE || + consumers.front()->inputs.size() != 1 || consumers.front()->inputs[0] != match.rows.output->id) { + return {}; + } + + const GraphNode * scale_node = consumers.front(); + const ScaleParams * params = op_params_as(scale_node->params); + const Value * output = graph_value(context.graph, scale_node->output); + if (params == nullptr || !std::isfinite(params->scale) || !std::isfinite(params->bias) || output == nullptr || + output->type != GGML_TYPE_F32 || !output->contiguous || output->alias_source.value >= 0 || + output->ne != match.rows.output->ne || output->byte_count != match.rows.output->byte_count) { + return {}; + } + + const std::array buffers = { match.rows.ids, match.rows.weight, match.rows.output, output }; + for (size_t lhs = 0; lhs < buffers.size(); ++lhs) { + for (size_t rhs = lhs + 1; rhs < buffers.size(); ++rhs) { + if (buffers[lhs]->storage_root == buffers[rhs]->storage_root) { + return {}; + } + } + } + if (!context.graph.index().node_index(scale_node, match.scale_index)) { + return {}; + } + + match.scale_node = scale_node; + match.output = output; + match.scale = params->scale; + match.bias = params->bias; + return match; +} + +static GetRowsRmsNormMatch match_get_rows_rmsnorm_binary(const DispatchMatchContext & context) { + GetRowsRmsNormMatch match; + match.rows = match_get_rows_f32(context.graph, context.root_node); + if (!match.rows.matched() || !context.graph.has_index() || match.rows.hidden_size < 128 || + match.rows.hidden_size > 32768 || match.rows.hidden_size % 128 != 0) { + return {}; + } + + const GraphNode * rms = nullptr; + for (const GraphNode * consumer : context.graph.index().consumers(match.rows.output->id)) { + if (consumer != nullptr && consumer->op == GGML_OP_RMS_NORM && consumer->inputs.size() == 1 && + consumer->inputs[0] == match.rows.output->id) { + if (rms != nullptr) { + return {}; + } + rms = consumer; + } + } + if (rms == nullptr) { + return {}; + } + + const RmsNormParams * params = op_params_as(rms->params); + const Value * rms_output = graph_value(context.graph, rms->output); + if (params == nullptr || !std::isfinite(params->eps) || params->eps <= 0.0f || rms_output == nullptr || + rms_output->type != GGML_TYPE_F32 || !rms_output->contiguous || rms_output->ne != match.rows.output->ne) { + return {}; + } + + const std::vector & rms_consumers = context.graph.index().consumers(rms_output->id); + if (rms_consumers.size() != 1 || rms_consumers.front() == nullptr || rms_consumers.front()->op != GGML_OP_MUL || + rms_consumers.front()->inputs.size() != 2) { + return {}; + } + const GraphNode * binary = rms_consumers.front(); + const ValueId weight_id = binary->inputs[0] == rms_output->id ? binary->inputs[1] : + binary->inputs[1] == rms_output->id ? binary->inputs[0] : + ValueId(); + const Value * norm_weight = graph_value(context.graph, weight_id); + const Value * output = graph_value(context.graph, binary->output); + if (weight_id.value < 0 || norm_weight == nullptr || output == nullptr || norm_weight->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !norm_weight->contiguous || !output->contiguous || + norm_weight->ne[0] != match.rows.hidden_size || norm_weight->ne[1] != 1 || norm_weight->ne[2] != 1 || + norm_weight->ne[3] != 1 || output->ne != match.rows.output->ne || + match.rows.output->storage_root == output->storage_root || norm_weight->storage_root == output->storage_root) { + return {}; + } + + if (!context.graph.index().node_index(rms, match.rms_node_index) || + !context.graph.index().node_index(binary, match.binary_node_index)) { + return {}; + } + match.rms_node = rms; + match.binary_node = binary; + match.norm_weight = norm_weight; + match.normalized_output = output; + match.epsilon = params->eps; + return match; +} + +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); + 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; +} + +static Dispatch make_get_rows_next_dispatch(const GetRowsMatch & match, + CommonMulMatWeightFormat next_format, + ValueId next_value, + size_t next_byte_count) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGetRowsF32NextKernel); + add_common_integer_parameters(dispatch, match); + add_common_compile_parameters(dispatch, match); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_f32.next_format", + to_config_value(common_mul_mat_format_config_value(next_format))); + add_primary_bindings(dispatch, match); + dispatch.bindings.push_back({ next_value, 0, next_byte_count }); + return dispatch; +} + +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()) { + return false; + } + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(make_get_rows_dispatch(match)); + return true; +} + +static bool match_get_rows_scale_f32_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const GetRowsScaleMatch match = match_get_rows_scale_f32(context); + if (!match.matched() || match.scale_index >= context.covered_nodes.size() || + context.covered_nodes[context.root_index] || context.covered_nodes[match.scale_index]) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGetRowsScaleF32Kernel); + add_common_integer_parameters(dispatch, match.rows); + add_common_compile_parameters(dispatch, match.rows); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_scale_f32.scale", to_config_value(match.scale)); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_scale_f32.bias", to_config_value(match.bias)); + dispatch.bindings.push_back({ match.rows.ids->id, 0, match.rows.ids->byte_count }); + dispatch.bindings.push_back({ match.rows.weight->id, 0, match.rows.weight->byte_count }); + dispatch.bindings.push_back({ match.rows.output->id, 0, match.rows.output->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.covered_nodes.push_back(match.scale_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_get_rows_f32_next_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const GetRowsMatch match = match_get_rows_f32(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + 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(); + + CommonMulMatWeightFormat next_format; + if (!common_mul_mat_alternate_format_for_type(GGML_TYPE_Q8_1, next_format)) { + return false; + } + 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(); +} + +static bool match_get_rows_rmsnorm_binary_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const GetRowsRmsNormMatch match = match_get_rows_rmsnorm_binary(context); + if (!match.matched() || match.rms_node_index >= context.covered_nodes.size() || + match.binary_node_index >= context.covered_nodes.size() || context.covered_nodes[context.root_index] || + context.covered_nodes[match.rms_node_index] || context.covered_nodes[match.binary_node_index]) { + return false; + } + + 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; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kGetRowsRmsNormBinaryQ8_1X4F16Kernel); + add_common_integer_parameters(dispatch, match.rows); + add_common_compile_parameters(dispatch, match.rows); + dispatch.kernel.compile_parameters.emplace("ggml.get_rows_rmsnorm.rms_epsilon", to_config_value(match.epsilon)); + dispatch.bindings.push_back({ match.rows.ids->id, 0, match.rows.ids->byte_count }); + dispatch.bindings.push_back({ match.rows.weight->id, 0, match.rows.weight->byte_count }); + 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_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.dispatches.push_back(std::move(dispatch)); + return dispatch_match.status.success(); +} + +} // namespace + +void register_get_rows_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.get_rows_scale.f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_get_rows_scale_f32_dispatch, + }); + registry.add({ + "common.get_rows_rmsnorm_binary.q8_1_x4_f16", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 330, + DispatchSource::Common, + match_get_rows_rmsnorm_binary_dispatch, + }); + registry.add({ + "common.get_rows.f32_next", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 150, + DispatchSource::Common, + match_get_rows_f32_next_dispatch, + }); + registry.add({ + "common.get_rows.f32", + GGML_OP_GET_ROWS, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Common, + match_get_rows_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.h new file mode 100644 index 000000000000..579f725c3ed0 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-get-rows.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch_registration/dispatch-registry.h" + +namespace ggml::hrx { + +void register_get_rows_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp new file mode 100644 index 000000000000..66f13832a181 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.cpp @@ -0,0 +1,200 @@ +#include "dispatch-glu.h" + +#include "dispatch-layout-utils.h" +#include "dispatch-mul-mat-common.h" +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kBinaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_binary_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool positive_shape(const Value & value) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] <= 0) { + return false; + } + } + return value.element_count > 0; +} + +static bool packed_f32_layout(const Value & value) { + size_t expected_stride = sizeof(float); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.nb[i] != expected_stride) { + return false; + } + expected_stride *= static_cast(value.ne[i]); + } + return true; +} + +static bool storage_ranges_disjoint(const Value & lhs, + size_t lhs_offset, + size_t lhs_byte_count, + const Value & rhs, + size_t rhs_offset, + size_t rhs_byte_count) { + if (lhs.storage != rhs.storage) { + return true; + } + if (lhs.storage_offset > std::numeric_limits::max() - lhs_offset || + rhs.storage_offset > std::numeric_limits::max() - rhs_offset) { + return false; + } + const size_t lhs_start = lhs.storage_offset + lhs_offset; + const size_t rhs_start = rhs.storage_offset + rhs_offset; + if (lhs_start > std::numeric_limits::max() - lhs_byte_count || + rhs_start > std::numeric_limits::max() - rhs_byte_count) { + return false; + } + return lhs_start + lhs_byte_count <= rhs_start || rhs_start + rhs_byte_count <= lhs_start; +} + +static bool packed_glu_half_span_bytes(const Value & input, + const Value & output, + size_t half_offset, + size_t & byte_count) { + if (input.type != GGML_TYPE_F32 || input.nb[0] != sizeof(float)) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (input.nb[i] % sizeof(float) != 0) { + return false; + } + } + + size_t max_offset = 0; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (output.ne[i] <= 0) { + return false; + } + const size_t extent = static_cast(output.ne[i] - 1); + if (extent != 0 && input.nb[i] > std::numeric_limits::max() / extent) { + return false; + } + const size_t dim_offset = extent * input.nb[i]; + if (max_offset > std::numeric_limits::max() - dim_offset) { + return false; + } + max_offset += dim_offset; + } + if (max_offset > std::numeric_limits::max() - sizeof(float)) { + return false; + } + byte_count = max_offset + sizeof(float); + if (input.storage_offset > std::numeric_limits::max() - half_offset || + input.storage_offset + half_offset > std::numeric_limits::max() - byte_count) { + return false; + } + return input.storage_offset + half_offset + byte_count <= input.storage_byte_count; +} + +static void add_packed_glu_parameters(Dispatch & dispatch, + const Value & input, + const Value & output, + size_t half_span) { + dispatch.kernel.integer_parameters.emplace("element_count", output.element_count); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(output.ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", std::to_string(output.ne[1])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", std::to_string(output.ne[2])); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", + std::to_string(input.nb[1] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", + std::to_string(input.nb[2] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", + std::to_string(input.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", + std::to_string(input.nb[1] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", + std::to_string(input.nb[2] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride3", + std::to_string(input.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_span", std::to_string(half_span / sizeof(float))); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_span", std::to_string(half_span / sizeof(float))); +} + +static void bind_packed_glu_source(Dispatch & dispatch, const Value & input, size_t offset, size_t byte_count) { + dispatch.bindings.push_back({ input.storage_root, input.storage_offset + offset, byte_count }); +} + +static bool match_packed_glu_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_GLU || node->inputs.size() != 1) { + return false; + } + + BinaryKind op; + if (!common_fused_binary_kind_from_params(node->params, op)) { + return false; + } + + const GluParams * params = op_params_as(node->params); + if (params == nullptr) { + return false; + } + + const Value * input = graph_value(context.graph, node->inputs[0]); + const Value * output = graph_value(context.graph, node->output); + if (input == nullptr || output == nullptr || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !positive_shape(*output) || input->ne[0] != 2 * output->ne[0] || output->alias_source.value >= 0 || + !output->contiguous || !packed_f32_layout(*output) || + static_cast(output->element_count) > std::numeric_limits::max()) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (input->ne[i] != output->ne[i]) { + return false; + } + } + + const size_t half_offset = static_cast(output->ne[0]) * sizeof(float); + size_t half_span = 0; + size_t input_span = 0; + if (!packed_glu_half_span_bytes(*input, *output, 0, half_span) || + !packed_glu_half_span_bytes(*input, *output, half_offset, half_span) || + !strided_f32_storage_span_bytes(*input, input_span) || + !storage_ranges_disjoint(*input, 0, input_span, *output, 0, output->byte_count)) { + return false; + } + + const size_t gate_offset = params->swapped ? half_offset : 0; + const size_t up_offset = params->swapped ? 0 : half_offset; + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kBinaryF32Kernel); + add_packed_glu_parameters(dispatch, *input, *output, half_span); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.op", std::to_string(binary_kind_config_value(op))); + bind_packed_glu_source(dispatch, *input, gate_offset, half_span); + bind_packed_glu_source(dispatch, *input, up_offset, half_span); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +} // namespace + +void register_glu_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.packed_glu_f32", + GGML_OP_GLU, + DispatchMatchKind::SingleOp, + 10, + DispatchSource::Common, + match_packed_glu_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.h new file mode 100644 index 000000000000..fec5fa69ed07 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-glu.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_glu_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h new file mode 100644 index 000000000000..e0c6d594732a --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-layout-utils.h @@ -0,0 +1,56 @@ +#pragma once + +#include "ggml.h" +#include "graph/value-map.h" + +#include +#include + +namespace ggml::hrx { + +inline bool strided_f32_storage_span_bytes(const Value & value, size_t & byte_count) { + if (value.type != GGML_TYPE_F32 || value.nb[0] != sizeof(float)) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (value.nb[i] % sizeof(float) != 0) { + return false; + } + } + + size_t max_offset = 0; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] <= 0) { + return false; + } + const size_t extent = static_cast(value.ne[i] - 1); + if (extent != 0 && value.nb[i] > std::numeric_limits::max() / extent) { + return false; + } + const size_t dim_offset = extent * value.nb[i]; + if (max_offset > std::numeric_limits::max() - dim_offset) { + return false; + } + max_offset += dim_offset; + } + if (max_offset > std::numeric_limits::max() - sizeof(float)) { + return false; + } + + byte_count = max_offset + sizeof(float); + if (value.storage_offset > std::numeric_limits::max() - byte_count) { + return false; + } + return value.storage_offset + byte_count <= value.storage_byte_count; +} + +inline bool strided_f32_storage_span_elements(const Value & value, size_t & element_count) { + size_t byte_count = 0; + if (!strided_f32_storage_span_bytes(value, byte_count) || byte_count % sizeof(float) != 0) { + return false; + } + element_count = byte_count / sizeof(float); + return true; +} + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..93127a1e6c58 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h @@ -0,0 +1,662 @@ +#pragma once + +#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" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +struct CommonMulMatMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + KernelCatalogRef kernel = {}; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + CommonMulMatWeightFormat weight_format = CommonMulMatWeightFormat::Q4K; + UnaryKind output_unary_op = UnaryKind::Identity; + size_t unary_node_index = 0; + bool has_fused_unary = false; + + bool matched() const { + return input != nullptr && weight != nullptr && output != nullptr && kernel.id != kUncatalogedKernelId; + } +}; + +struct CommonSymmetricI4ActivationLayout { + size_t payload_bytes = 0; + size_t scales_offset = 0; + size_t metadata_bytes = 0; + size_t sums_offset = 0; + size_t total_bytes = 0; +}; + +inline constexpr const char kCommonSymmetricI4K32ActivationAlternateName[] = + "common.mul_mat.symmetric_i4_k32.activation"; + +enum class CommonQ8ActivationPolicy { + ExistingAlternateOnly, + AllowStandaloneQuantize, +}; + +inline size_t common_align_up(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +inline CommonSymmetricI4ActivationLayout common_symmetric_i4_activation_layout(int64_t input_size, + int64_t token_count) { + const size_t element_count = static_cast(input_size) * static_cast(token_count); + const size_t payload_bytes = element_count / 2; + const size_t metadata_bytes = element_count / 8; + const size_t scales_offset = common_align_up(payload_bytes, 256); + const size_t sums_offset = common_align_up(scales_offset + metadata_bytes, 256); + return { + payload_bytes, scales_offset, metadata_bytes, sums_offset, common_align_up(sums_offset + metadata_bytes, 256), + }; +} + +inline size_t common_symmetric_shared4_row_group_size(int64_t input_size, int64_t output_size, int64_t quant_bits) { + const size_t logical_output_size = static_cast(output_size); + const size_t materialized_row_bytes = 4 + 8 * 32 * static_cast(quant_bits) / 8; + const size_t materialized_bytes = + logical_output_size * static_cast(input_size / 256) * materialized_row_bytes; + if (quant_bits == 4 && materialized_bytes < size_t{ 16 } * 1024 * 1024) { + return 32; + } + if (quant_bits == 4 && materialized_bytes <= size_t{ 32 } * 1024 * 1024 && output_size > input_size) { + return 96; + } + return ((logical_output_size + 255) / 256) * 32; +} + +inline size_t common_symmetric_shared4_weight_byte_count(int64_t input_size, int64_t output_size, int64_t quant_bits) { + const size_t logical_output_size = static_cast(output_size); + const size_t row_group_size = common_symmetric_shared4_row_group_size(input_size, output_size, quant_bits); + const size_t padded_output_size = (logical_output_size + row_group_size - 1) / row_group_size * row_group_size; + const size_t materialized_row_bytes = 4 + 8 * 32 * static_cast(quant_bits) / 8; + return padded_output_size * static_cast(input_size / 256) * materialized_row_bytes; +} + +inline size_t common_symmetric_i4_shared4_weight_byte_count(int64_t input_size, int64_t output_size) { + return common_symmetric_shared4_weight_byte_count(input_size, output_size, 4); +} + +inline DispatchBinding common_symmetric_i4_shared4_weight_binding(const Value & weight, + int64_t input_size, + int64_t output_size) { + DispatchBinding binding; + binding.value = weight.id; + binding.length = common_symmetric_i4_shared4_weight_byte_count(input_size, output_size); + binding.layout = kSymmetricI4K32EightGroupsShared4Layout; + binding.source_type = weight.type; + binding.input_size = input_size; + binding.output_size = output_size; + binding.source_length = weight.byte_count; + return binding; +} + +inline DispatchBinding common_symmetric_i4_shared4_multistart_weight_binding(const Value & weight, + int64_t input_size, + int64_t output_size) { + DispatchBinding binding = common_symmetric_i4_shared4_weight_binding(weight, input_size, output_size); + binding.layout = kSymmetricI4K32EightGroupsShared4MultistartLayout; + return binding; +} + +inline const Value * common_graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +inline bool common_same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +inline bool common_is_supported_dense_input_size(int64_t input_size) { + return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; +} + +inline bool common_is_supported_dense_input_size(CommonMulMatWeightFormat format, int64_t input_size) { + return common_mul_mat_supported_dense_input_size(format, input_size); +} + +inline bool common_is_supported_dense_output_size(int64_t output_size) { + return output_size >= 1 && output_size <= 262144; +} + +// Select the packed prefill schedule. Long Q4 contractions only benefit when +// their producer supplies the layout without a separate conversion pass. +inline bool common_mul_mat_uses_k16_major_f16(CommonMulMatWeightFormat format, + int64_t input_size, + int64_t output_size, + int64_t token_count, + bool packed_producer = false) { + if (token_count < 512 || token_count > 2048 || token_count % 512 != 0 || input_size % 256 != 0 || + output_size % 64 != 0) { + return false; + } + if (format == CommonMulMatWeightFormat::Q4K) { + return token_count == 512 && output_size / 64 >= 64 && + (input_size <= 2 * output_size || (packed_producer && input_size <= 4 * output_size)); + } + if (format == CommonMulMatWeightFormat::Q6K) { + const int64_t tile = input_size > output_size ? 64 : 128; + return output_size % tile == 0 && (token_count / 512) * (output_size / tile) >= 32; + } + return false; +} + +inline ggml_type common_mul_mat_format_type(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::Q1_0: + return GGML_TYPE_Q1_0; + case CommonMulMatWeightFormat::Q3K: + return GGML_TYPE_Q3_K; + case CommonMulMatWeightFormat::Q4K: + case CommonMulMatWeightFormat::Q4KRow64: + return GGML_TYPE_Q4_K; + case CommonMulMatWeightFormat::Q5K: + return GGML_TYPE_Q5_K; + case CommonMulMatWeightFormat::Q6K: + case CommonMulMatWeightFormat::Q6KRow64: + return GGML_TYPE_Q6_K; + case CommonMulMatWeightFormat::Q4_0: + return GGML_TYPE_Q4_0; + case CommonMulMatWeightFormat::Q4_1: + return GGML_TYPE_Q4_1; + case CommonMulMatWeightFormat::Q5_0: + 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: + return GGML_TYPE_IQ3_S; + case CommonMulMatWeightFormat::IQ4_NL: + return GGML_TYPE_IQ4_NL; + case CommonMulMatWeightFormat::IQ4_XS: + return GGML_TYPE_IQ4_XS; + case CommonMulMatWeightFormat::Q8_0: + return GGML_TYPE_Q8_0; + case CommonMulMatWeightFormat::Q8_1: + return GGML_TYPE_Q8_1; + case CommonMulMatWeightFormat::F16: + return GGML_TYPE_F16; + case CommonMulMatWeightFormat::BF16: + return GGML_TYPE_BF16; + case CommonMulMatWeightFormat::F32: + return GGML_TYPE_F32; + } + return GGML_TYPE_COUNT; +} + +inline bool common_is_supported_dense_decode_output_size(CommonMulMatWeightFormat format, + int64_t input_size, + int64_t output_size) { + static constexpr uint64_t kMaxDenseDecodeOutputSize = 1048576; + static constexpr uint64_t kAmdgpuAddressableByteRangeSize = uint64_t{ 1 } << 32; + + if (output_size < 1) { + return false; + } + + const ggml_type type = common_mul_mat_format_type(format); + if (type == GGML_TYPE_COUNT) { + return false; + } + + const size_t weight_row_size = ggml_row_size(type, input_size); + if (weight_row_size == 0) { + return false; + } + + const uint64_t addressable_output_size = kAmdgpuAddressableByteRangeSize / static_cast(weight_row_size); + const uint64_t max_output_size = std::min(kMaxDenseDecodeOutputSize, addressable_output_size); + return static_cast(output_size) <= max_output_size; +} + +inline bool common_is_supported_rmsnorm_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +inline bool common_is_supported_prefill_token_count(int64_t token_count) { + return token_count > 1 && token_count <= 2048; +} + +inline bool common_is_supported_decode_token_count(int64_t token_count) { + return token_count == 1; +} + +inline bool common_is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +inline std::string common_to_config_value(int64_t value) { + return std::to_string(value); +} + +inline std::string common_to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +inline int64_t common_ceil_div(int64_t value, int64_t divisor) { + return (value + divisor - 1) / divisor; +} + +inline bool common_is_weight_shape(const Value & weight, int64_t hidden_size) { + if (weight.ne[0] != hidden_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (weight.ne[i] != 1) { + return false; + } + } + return true; +} + +inline const GraphNode * common_find_single_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + if (!graph.has_index()) { + return nullptr; + } + const GraphNode * match = nullptr; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer == nullptr || consumer->op != op) { + continue; + } + if (match != nullptr) { + return nullptr; + } + match = consumer; + } + return match; +} + +inline const GraphNode * common_find_only_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + if (!graph.has_index()) { + return nullptr; + } + const std::vector & consumers = graph.index().consumers(value); + if (consumers.size() != 1 || consumers.front() == nullptr || consumers.front()->op != op) { + return nullptr; + } + return consumers.front(); +} + +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 && common_symmetric_i4_lowrow_mul_mat_eligible(graph, *consumer, value)) { + return true; + } + } + + for (const GraphNode * consumer : graph.index().consumers(value.id)) { + if (consumer == nullptr || !is_layout_alias_node(graph, *consumer) || consumer->inputs.size() != 1 || + consumer->inputs.front() != value.id) { + continue; + } + + 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) { + 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; +} + +inline bool common_binary_node_is_mul(const GraphNode & node) { + if (node.op != GGML_OP_MUL || node.inputs.size() != 2) { + return false; + } + const BinaryParams * binary_params = op_params_as(node.params); + return binary_params == nullptr || binary_params->op == BinaryKind::Mul; +} + +inline bool common_binary_node_is_add(const GraphNode & node) { + if (node.op != GGML_OP_ADD || node.inputs.size() != 2) { + return false; + } + const BinaryParams * binary_params = op_params_as(node.params); + return binary_params == nullptr || binary_params->op == BinaryKind::Add; +} + +inline bool common_is_swiglu_params(const OpParams & params) { + const BinaryParams * binary_params = op_params_as(params); + if (binary_params != nullptr) { + return binary_params->op == BinaryKind::SwiGLU; + } + + const GluParams * glu_params = op_params_as(params); + return glu_params != nullptr && glu_params->op == GGML_GLU_OP_SWIGLU; +} + +inline CommonMulMatMatch common_match_mul_mat_any_format(const Graph & graph, + const GraphNode * node, + KernelCatalogRef kernel, + bool decode) { + CommonMulMatMatch match; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return match; + } + + const Value * weight = common_graph_value(graph, node->inputs[0]); + const Value * input = common_graph_value(graph, node->inputs[1]); + const Value * output = common_graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr || !common_is_2d(*weight) || !common_is_2d(*input) || + !common_is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || + input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return {}; + } + + 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 {}; + } + + const int64_t input_size = weight->ne[0]; + const int64_t output_size = weight->ne[1]; + const int64_t token_count = input->ne[1]; + const bool token_count_supported = decode ? common_is_supported_decode_token_count(token_count) : + common_is_supported_prefill_token_count(token_count); + if (input->ne[0] != input_size || output->ne[0] != output_size || output->ne[1] != token_count || + !token_count_supported || !common_is_supported_dense_input_size(format, input_size) || + !(decode ? common_is_supported_dense_decode_output_size(format, input_size, output_size) : + common_is_supported_dense_output_size(output_size))) { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.kernel = kernel; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + match.weight_format = format; + 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, + int64_t input_size, + int64_t token_count, + DispatchMatch & match, + DispatchBinding & binding) { + const size_t bytes = static_cast(token_count * input_size) * sizeof(ggml_fp16_t); + const CommandPlanAlternateValue * alternate = + find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_F16, bytes); + if (alternate != nullptr) { + binding = { alternate->alternate_value, 0, bytes }; + return true; + } + + constexpr KernelCatalogRef kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_f32_f16"); + constexpr const char * name = "common.mul_mat.f16"; + const ValueId activation = context.next_plan_value; + match.transients.push_back({ activation, name, bytes, 256 }); + Dispatch convert; + convert.kernel = make_kernel_specialization(kernel); + convert.kernel.integer_parameters.emplace("element_count", token_count * input_size); + convert.bindings.push_back({ input.id, 0, input.byte_count }); + convert.bindings.push_back({ activation, 0, bytes }); + match.dispatches.push_back(std::move(convert)); + Status status; + if (!match.metadata.append_alternate_value({ input.id, activation, GGML_TYPE_F16, bytes, name }, status)) { + match.status.append(status); + return false; + } + binding = { activation, 0, bytes }; + return true; +} + +// Private K16-major copy; it is not published as an ordinary F16 alternate. +inline bool common_prepare_k16_major_f16_input(const DispatchMatchContext & context, + const Value & input, + int64_t input_size, + int64_t token_count, + DispatchMatch & match, + DispatchBinding & binding) { + constexpr KernelCatalogRef kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_f16_k16_major"); + constexpr const char * name = "common.mul_mat.k16_major_f16"; + const size_t bytes = static_cast(token_count * input_size) * sizeof(ggml_fp16_t); + const CommandPlanGeneratedResource * generated = + context.plan.metadata.find_generated_resource(input.id, GeneratedResourceRole::F16K16Major); + if (generated != nullptr && generated->byte_count == bytes) { + binding = { generated->generated_value, 0, bytes }; + return true; + } + const CommandPlanAlternateValue * alternate = + find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_F16, bytes); + const ValueId packed(context.next_plan_value.value + static_cast(match.transients.size())); + match.transients.push_back({ packed, name, bytes, 256 }); + Dispatch copy; + copy.kernel = make_kernel_specialization(kernel); + copy.kernel.integer_parameters.emplace("token_count", token_count); + copy.kernel.compile_parameters.emplace("ggml.copy_f16_k16_major.input_size", common_to_config_value(input_size)); + copy.kernel.compile_parameters.emplace("ggml.copy_f16_k16_major.token_count", common_to_config_value(token_count)); + copy.kernel.compile_parameters.emplace("ggml.copy_f16_k16_major.input_is_f16", alternate != nullptr ? "1" : "0"); + copy.bindings.push_back(alternate != nullptr ? DispatchBinding{ alternate->alternate_value, 0, bytes } : + DispatchBinding{ input.id, 0, input.byte_count }); + copy.bindings.push_back({ packed, 0, bytes }); + match.dispatches.push_back(std::move(copy)); + Status status; + if (!match.metadata.append_generated_resource({ input.id, GeneratedResourceRole::F16K16Major, packed, bytes, {} }, + status)) { + match.status.append(status); + return false; + } + binding = { packed, 0, bytes }; + return true; +} + +// Share one activation producer between ordinary and fused matmuls. +inline bool common_prepare_q8_1_x4_input(const DispatchMatchContext & context, + const Value & input, + int64_t input_size, + int64_t token_count, + DispatchMatch & match, + DispatchBinding & binding, + CommonQ8ActivationPolicy policy) { + const size_t bytes = static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, input_size); + const CommandPlanAlternateValue * alternate = + find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_Q8_1, bytes); + if (alternate != nullptr) { + binding = { alternate->alternate_value, 0, bytes }; + return true; + } + if (policy != CommonQ8ActivationPolicy::AllowStandaloneQuantize) { + return false; + } + + constexpr KernelCatalogRef kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_quantize_q8_1_x4_f32"); + constexpr const char * name = "common.mul_mat.q8_1_x4"; + const ValueId activation = context.next_plan_value; + match.transients.push_back({ activation, name, bytes, 256 }); + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kernel); + quantize.kernel.integer_parameters.emplace("token_count", token_count); + quantize.kernel.integer_parameters.emplace("input_size", input_size); + quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + common_to_config_value(token_count * input_size / 128)); + quantize.bindings.push_back({ input.id, 0, input.byte_count }); + quantize.bindings.push_back({ activation, 0, bytes }); + match.dispatches.push_back(std::move(quantize)); + + Status status; + if (!match.metadata.append_alternate_value({ input.id, activation, GGML_TYPE_Q8_1, bytes, name }, status)) { + match.status.append(status); + return false; + } + binding = { activation, 0, bytes }; + return true; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h new file mode 100644 index 000000000000..843956d9c136 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id-common.h @@ -0,0 +1,352 @@ +#pragma once + +#include "../dispatch-registry.h" +#include "dispatch-mul-mat-weight-format.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +struct CommonMulMatIdMatch { + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * output = nullptr; + const Value * route_ids = nullptr; + CommandPlanMoeRoutingBundle routing_bundle; + bool has_routing_bundle = false; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + int64_t route_count = 0; + int64_t input_route_count = 0; + int64_t expert_count = 0; + CommonMulMatWeightFormat weight_format = CommonMulMatWeightFormat::Q4K; + + bool matched() const { return input != nullptr && weight != nullptr && output != nullptr && route_ids != nullptr; } +}; + +inline const Value * common_mul_mat_id_graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +inline bool common_mul_mat_id_is_shape(const Value & value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == ne2 && value.ne[3] == ne3; +} + +inline bool common_mul_mat_id_same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +inline bool common_mul_mat_id_supported_dense_input_size(CommonMulMatWeightFormat format, int64_t input_size) { + return common_mul_mat_supported_dense_input_size(format, input_size); +} + +inline bool common_mul_mat_id_supported_dense_output_size(int64_t output_size) { + return output_size >= 1 && output_size <= 262144; +} + +inline bool common_mul_mat_id_supported_rmsnorm_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +inline bool common_mul_mat_id_supported_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= 2048; +} + +inline bool common_mul_mat_id_supported_route_count(int64_t route_count) { + return route_count >= 1 && route_count <= 32; +} + +inline bool common_mul_mat_id_supported_expert_count(int64_t expert_count) { + return expert_count >= 1 && expert_count <= 512; +} + +inline std::string common_mul_mat_id_to_config_value(int64_t value) { + return std::to_string(value); +} + +inline std::string common_mul_mat_id_to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +inline int64_t common_mul_mat_id_ceil_div(int64_t value, int64_t divisor) { + return (value + divisor - 1) / divisor; +} + +inline int64_t common_mul_mat_id_partition_descriptor_capacity(int64_t token_count, + int64_t route_count, + int64_t expert_count) { + return common_mul_mat_id_ceil_div(token_count * route_count, 32) + expert_count; +} + +inline const GraphNode * common_mul_mat_id_find_single_consumer_with_op(const Graph & graph, + ValueId value, + ggml_op op) { + if (!graph.has_index()) { + return nullptr; + } + const GraphNode * match = nullptr; + for (const GraphNode * consumer : graph.index().consumers(value)) { + if (consumer == nullptr || consumer->op != op) { + continue; + } + if (match != nullptr) { + return nullptr; + } + match = consumer; + } + return match; +} + +inline const GraphNode * common_mul_mat_id_find_only_consumer_with_op(const Graph & graph, ValueId value, ggml_op op) { + if (!graph.has_index()) { + return nullptr; + } + const std::vector & consumers = graph.index().consumers(value); + if (consumers.size() != 1 || consumers.front() == nullptr || consumers.front()->op != op) { + return nullptr; + } + return consumers.front(); +} + +inline bool common_mul_mat_id_binary_node_is_mul(const GraphNode & node) { + if (node.op != GGML_OP_MUL || node.inputs.size() != 2) { + return false; + } + const BinaryParams * binary_params = op_params_as(node.params); + return binary_params == nullptr || binary_params->op == BinaryKind::Mul; +} + +inline bool common_mul_mat_id_binary_node_is_add(const GraphNode & node) { + if (node.op != GGML_OP_ADD || node.inputs.size() != 2) { + return false; + } + const BinaryParams * binary_params = op_params_as(node.params); + return binary_params == nullptr || binary_params->op == BinaryKind::Add; +} + +inline bool common_mul_mat_id_is_swiglu_params(const OpParams & params) { + const BinaryParams * binary_params = op_params_as(params); + if (binary_params != nullptr) { + return binary_params->op == BinaryKind::SwiGLU; + } + + const GluParams * glu_params = op_params_as(params); + return glu_params != nullptr && glu_params->op == GGML_GLU_OP_SWIGLU; +} + +inline bool common_mul_mat_id_is_weight_shape(const Value & weight, int64_t hidden_size) { + if (weight.ne[0] != hidden_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (weight.ne[i] != 1) { + return false; + } + } + return true; +} + +inline size_t common_mul_mat_id_expert_table_size(int64_t token_count, int64_t expert_count) { + return static_cast(expert_count + expert_count * token_count) * sizeof(int32_t); +} + +inline size_t common_mul_mat_id_partition_table_size(int64_t token_count, int64_t route_count, int64_t expert_count) { + const int64_t assignment_count = token_count * route_count; + const int64_t assignment_partition_count = (assignment_count + 31) / 32; + return static_cast(1 + assignment_partition_count + expert_count) * sizeof(int32_t); +} + +inline bool common_mul_mat_id_bundle_matches(const CommandPlanMoeRoutingBundle & bundle, + ValueId route_ids, + int64_t token_count, + int64_t route_count, + int64_t expert_count) { + return bundle.route_ids == route_ids && bundle.expert_table.value >= 0 && bundle.partition_table.value >= 0 && + bundle.expert_table_byte_count == common_mul_mat_id_expert_table_size(token_count, expert_count) && + bundle.partition_table_byte_count == + common_mul_mat_id_partition_table_size(token_count, route_count, expert_count) && + bundle.token_count == token_count && bundle.route_count == route_count && + bundle.expert_count == expert_count && bundle.route_stride >= route_count; +} + +inline CommonMulMatIdMatch common_match_mul_mat_id_any_format(const Graph & graph, + const GraphNode * node, + const CommandPlan & plan) { + CommonMulMatIdMatch match; + if (node == nullptr || node->op != GGML_OP_MUL_MAT_ID || node->inputs.size() != 3 || !graph.has_index()) { + return match; + } + + const Value * weight = common_mul_mat_id_graph_value(graph, node->inputs[0]); + const Value * input = common_mul_mat_id_graph_value(graph, node->inputs[1]); + const Value * route_ids = common_mul_mat_id_graph_value(graph, node->inputs[2]); + const Value * output = common_mul_mat_id_graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || route_ids == nullptr || output == nullptr || !weight->contiguous || + !input->contiguous || !output->contiguous || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + route_ids->type != GGML_TYPE_I32) { + return {}; + } + + CommonMulMatWeightFormat format = CommonMulMatWeightFormat::Q4K; + 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 expert_count = weight->ne[2]; + const int64_t input_route_count = input->ne[1]; + const int64_t token_count = input->ne[2]; + const int64_t route_count = route_ids->ne[0]; + if (!common_mul_mat_id_is_shape(*weight, input_size, output_size, expert_count, 1) || + !common_mul_mat_id_is_shape(*input, input_size, input_route_count, token_count, 1) || + !common_mul_mat_id_is_shape(*route_ids, route_count, token_count, 1, 1) || + !common_mul_mat_id_is_shape(*output, output_size, route_count, token_count, 1) || + !common_mul_mat_id_supported_dense_input_size(format, input_size) || + !common_mul_mat_id_supported_dense_output_size(output_size) || + !common_mul_mat_id_supported_token_count(token_count) || + !common_mul_mat_id_supported_route_count(route_count) || + !common_mul_mat_id_supported_expert_count(expert_count) || input_route_count <= 0 || + input_route_count > route_count || route_count % input_route_count != 0) { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.route_ids = route_ids; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + match.route_count = route_count; + match.input_route_count = input_route_count; + match.expert_count = expert_count; + match.weight_format = format; + + const CommandPlanMoeRoutingBundle * routing_bundle = plan.metadata.find_moe_routing_bundle(route_ids->id); + if (routing_bundle != nullptr && + common_mul_mat_id_bundle_matches(*routing_bundle, route_ids->id, token_count, route_count, expert_count)) { + match.routing_bundle = *routing_bundle; + match.has_routing_bundle = true; + } + return match; +} + +inline void common_mul_mat_id_add_moe_routing_compile_parameters(Dispatch & dispatch, + const CommonMulMatIdMatch & match) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_mul_mat_id_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.route_count", + common_mul_mat_id_to_config_value(match.route_count)); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.expert_count", + common_mul_mat_id_to_config_value(match.expert_count)); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.descriptor_expert_mask", "511"); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.descriptor_partition_shift", "9"); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.descriptor_row_count_shift", "15"); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.partition_workgroup_size", "512"); +} + +inline bool common_mul_mat_id_ensure_moe_routing_bundle(const DispatchMatchContext & context, + const CommonMulMatIdMatch & match, + DispatchMatch & dispatch_match, + CommandPlanMoeRoutingBundle & routing_bundle) { + static constexpr KernelCatalogRef kMoeBuildExpertTableKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_moe_build_expert_table"); + static constexpr KernelCatalogRef kMoeBuildExpertPartitionTableKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_moe_build_expert_partition_table"); + + if (match.has_routing_bundle) { + routing_bundle = match.routing_bundle; + return true; + } + + const ValueId expert_table_value(context.next_plan_value.value + + static_cast(dispatch_match.transients.size())); + const ValueId partition_table_value(expert_table_value.value + 1); + const size_t expert_table_bytes = common_mul_mat_id_expert_table_size(match.token_count, match.expert_count); + const size_t partition_table_bytes = + common_mul_mat_id_partition_table_size(match.token_count, match.route_count, match.expert_count); + const int64_t route_stride = match.route_count; + const size_t route_ids_length = static_cast(match.token_count * route_stride) * sizeof(int32_t); + + dispatch_match.transients.push_back( + { expert_table_value, "common.moe_routing.expert_table", expert_table_bytes, 256 }); + dispatch_match.transients.push_back( + { partition_table_value, "common.moe_routing.partition_table", partition_table_bytes, 256 }); + + const CommandPlanResourceMetadata routing_metadata = make_command_plan_resource_metadata(MoeRoutingResourceMetadata{ + match.token_count, + match.route_count, + route_stride, + match.expert_count, + }); + routing_bundle = { + match.route_ids->id, ValueId(), expert_table_value, partition_table_value, expert_table_bytes, + partition_table_bytes, match.token_count, match.route_count, route_stride, match.expert_count, + }; + + Status metadata_status; + if (!dispatch_match.metadata.append_generated_resource( + { + match.route_ids->id, + GeneratedResourceRole::MoeExpertTable, + expert_table_value, + expert_table_bytes, + routing_metadata, + }, + metadata_status) || + !dispatch_match.metadata.append_generated_resource( + { + match.route_ids->id, + GeneratedResourceRole::MoePartitionTable, + partition_table_value, + partition_table_bytes, + routing_metadata, + }, + metadata_status) || + !dispatch_match.metadata.append_moe_routing_bundle(routing_bundle, metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + + Dispatch expert_table_dispatch; + expert_table_dispatch.kernel = make_kernel_specialization(kMoeBuildExpertTableKernel); + expert_table_dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + expert_table_dispatch.kernel.integer_parameters.emplace("route_count", match.route_count); + expert_table_dispatch.kernel.integer_parameters.emplace("route_stride", route_stride); + expert_table_dispatch.kernel.integer_parameters.emplace("expert_count", match.expert_count); + common_mul_mat_id_add_moe_routing_compile_parameters(expert_table_dispatch, match); + expert_table_dispatch.bindings.push_back({ match.route_ids->id, 0, route_ids_length }); + expert_table_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); + dispatch_match.dispatches.push_back(std::move(expert_table_dispatch)); + + Dispatch partition_table_dispatch; + partition_table_dispatch.kernel = make_kernel_specialization(kMoeBuildExpertPartitionTableKernel); + partition_table_dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + partition_table_dispatch.kernel.integer_parameters.emplace("route_count", match.route_count); + partition_table_dispatch.kernel.integer_parameters.emplace("expert_count", match.expert_count); + common_mul_mat_id_add_moe_routing_compile_parameters(partition_table_dispatch, match); + partition_table_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); + partition_table_dispatch.bindings.push_back({ partition_table_value, 0, partition_table_bytes }); + dispatch_match.dispatches.push_back(std::move(partition_table_dispatch)); + return true; +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp new file mode 100644 index 000000000000..5c9822c5aa9e --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.cpp @@ -0,0 +1,333 @@ +#include "dispatch-mul-mat-id.h" + +#include "dispatch-mul-mat-id-common.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kMulMatIdTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdTiledPostOpsF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnyPostOpsF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_input_f32_postops_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdTiledPostOpsNextRmsNormF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32"); +static constexpr KernelCatalogRef kMulMatIdSkinnyPostOpsNextRmsNormF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32"); + +struct MulMatIdPostOpsMatch { + CommonMulMatIdMatch root; + const Value * bias = nullptr; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const Value * norm_weight = nullptr; + const Value * normalized_output = nullptr; + std::vector add_nodes; + const GraphNode * rms_node = nullptr; + const GraphNode * mul_node = nullptr; + KernelCatalogRef kernel = {}; + float epsilon = 0.0f; + bool has_bias = false; + bool has_residual = false; + bool has_rmsnorm = false; + + bool matched() const { + return root.matched() && residual_output != nullptr && kernel.id != kUncatalogedKernelId && + (has_bias || has_residual); + } +}; + +static bool is_bias_shape(const Value & value, int64_t output_size) { + if (value.type != GGML_TYPE_F32 || !value.contiguous || value.ne[0] != output_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] != 1) { + return false; + } + } + return true; +} + +static bool common_mul_mat_id_is_skinny_route(const CommonMulMatIdMatch & match) { + return match.matched() && match.token_count <= 5; +} + +static KernelCatalogRef common_mul_mat_id_f32_kernel(const CommonMulMatIdMatch & match) { + return common_mul_mat_id_is_skinny_route(match) ? kMulMatIdSkinnyF32F32Kernel : kMulMatIdTiledF32F32Kernel; +} + +static KernelCatalogRef common_mul_mat_id_postops_f32_kernel(const CommonMulMatIdMatch & match, bool has_rmsnorm) { + if (has_rmsnorm) { + return common_mul_mat_id_is_skinny_route(match) ? kMulMatIdSkinnyPostOpsNextRmsNormF32F32Kernel : + kMulMatIdTiledPostOpsNextRmsNormF32F32Kernel; + } + return common_mul_mat_id_is_skinny_route(match) ? kMulMatIdSkinnyPostOpsF32F32Kernel : + kMulMatIdTiledPostOpsF32F32Kernel; +} + +static MulMatIdPostOpsMatch match_mul_mat_id_postops(const DispatchMatchContext & context) { + MulMatIdPostOpsMatch match; + const CommonMulMatIdMatch root = common_match_mul_mat_id_any_format(context.graph, context.root_node, context.plan); + if (!root.matched() || !context.graph.has_index()) { + return match; + } + + const Value * current = root.output; + for (int add_index = 0; add_index < 2; ++add_index) { + const GraphNode * add_node = + common_mul_mat_id_find_only_consumer_with_op(context.graph, current->id, GGML_OP_ADD); + if (add_node == nullptr || !common_mul_mat_id_binary_node_is_add(*add_node)) { + break; + } + + const bool current_is_lhs = add_node->inputs[0] == current->id; + const bool current_is_rhs = add_node->inputs[1] == current->id; + if (!current_is_lhs && !current_is_rhs) { + return {}; + } + + const ValueId other_id = current_is_lhs ? add_node->inputs[1] : add_node->inputs[0]; + const Value * other = common_mul_mat_id_graph_value(context.graph, other_id); + const Value * output = common_mul_mat_id_graph_value(context.graph, add_node->output); + if (other == nullptr || output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || + !common_mul_mat_id_same_shape(*root.output, *output)) { + return {}; + } + + if (!match.has_bias && is_bias_shape(*other, root.output_size)) { + match.bias = other; + match.add_nodes.push_back(add_node); + match.has_bias = true; + current = output; + continue; + } + + if (!match.has_residual && other->type == GGML_TYPE_F32 && other->contiguous && + common_mul_mat_id_same_shape(*root.output, *other)) { + match.residual_input = other; + match.add_nodes.push_back(add_node); + match.has_residual = true; + current = output; + continue; + } + + break; + } + + if (!match.has_bias && !match.has_residual) { + return {}; + } + + match.residual_output = current; + + const bool may_match_next_rmsnorm = + match.has_residual && common_mul_mat_id_supported_rmsnorm_hidden_size(root.output_size); + if (may_match_next_rmsnorm) { + const GraphNode * rms_node = + common_mul_mat_id_find_single_consumer_with_op(context.graph, current->id, GGML_OP_RMS_NORM); + if (rms_node != nullptr && rms_node->inputs.size() == 1) { + const RmsNormParams * rms_params = op_params_as(rms_node->params); + const Value * rms_output = common_mul_mat_id_graph_value(context.graph, rms_node->output); + if (rms_params != nullptr && std::isfinite(rms_params->eps) && rms_params->eps > 0.0f && + rms_output != nullptr && rms_output->type == GGML_TYPE_F32 && rms_output->contiguous && + common_mul_mat_id_same_shape(*current, *rms_output)) { + const GraphNode * mul_node = + common_mul_mat_id_find_single_consumer_with_op(context.graph, rms_node->output, GGML_OP_MUL); + if (mul_node != nullptr && common_mul_mat_id_binary_node_is_mul(*mul_node)) { + const bool rms_is_lhs = mul_node->inputs[0] == rms_node->output; + const bool rms_is_rhs = mul_node->inputs[1] == rms_node->output; + if (rms_is_lhs || rms_is_rhs) { + const ValueId norm_weight_id = rms_is_lhs ? mul_node->inputs[1] : mul_node->inputs[0]; + const Value * norm_weight = common_mul_mat_id_graph_value(context.graph, norm_weight_id); + const Value * normalized_output = + common_mul_mat_id_graph_value(context.graph, mul_node->output); + if (norm_weight != nullptr && normalized_output != nullptr && + norm_weight->type == GGML_TYPE_F32 && normalized_output->type == GGML_TYPE_F32 && + norm_weight->contiguous && normalized_output->contiguous && + common_mul_mat_id_is_weight_shape(*norm_weight, root.output_size) && + common_mul_mat_id_same_shape(*current, *normalized_output)) { + match.norm_weight = norm_weight; + match.normalized_output = normalized_output; + match.rms_node = rms_node; + match.mul_node = mul_node; + match.epsilon = rms_params->eps; + match.has_rmsnorm = true; + } + } + } + } + } + } + + match.root = root; + match.kernel = common_mul_mat_id_postops_f32_kernel(root, match.has_rmsnorm); + return match; +} + +} // namespace + +static bool match_mul_mat_id_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatIdMatch match = + common_match_mul_mat_id_any_format(context.graph, context.root_node, context.plan); + if (!match.matched()) { + return false; + } + + CommandPlanMoeRoutingBundle routing_bundle; + if (!common_mul_mat_id_ensure_moe_routing_bundle(context, match, dispatch_match, routing_bundle)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(common_mul_mat_id_f32_kernel(match)); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_mul_mat_id_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.input_size", + common_mul_mat_id_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.output_size", + common_mul_mat_id_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.expert_count", + common_mul_mat_id_to_config_value(match.expert_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.route_count", + common_mul_mat_id_to_config_value(match.route_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.input_route_count", + common_mul_mat_id_to_config_value(match.input_route_count)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_id.weight_format", + common_mul_mat_id_to_config_value(common_mul_mat_format_config_value(match.weight_format))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ routing_bundle.expert_table, 0, routing_bundle.expert_table_byte_count }); + dispatch.bindings.push_back({ routing_bundle.partition_table, 0, routing_bundle.partition_table_byte_count }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_mul_mat_id_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatIdPostOpsMatch match = match_mul_mat_id_postops(context); + if (!match.matched()) { + return false; + } + + CommandPlanMoeRoutingBundle routing_bundle; + if (!common_mul_mat_id_ensure_moe_routing_bundle(context, match.root, dispatch_match, routing_bundle)) { + return false; + } + + const int64_t completion_counter_count = + match.has_rmsnorm ? common_mul_mat_id_partition_descriptor_capacity( + match.root.token_count, match.root.route_count, match.root.expert_count) : + 0; + const ValueId completion_counters(context.next_plan_value.value + + static_cast(dispatch_match.transients.size())); + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.root.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_mul_mat_id_to_config_value(match.root.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.input_size", + common_mul_mat_id_to_config_value(match.root.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.output_size", + common_mul_mat_id_to_config_value(match.root.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.expert_count", + common_mul_mat_id_to_config_value(match.root.expert_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.route_count", + common_mul_mat_id_to_config_value(match.root.route_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.input_route_count", + common_mul_mat_id_to_config_value(match.root.input_route_count)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_id_postops.weight_format", + common_mul_mat_id_to_config_value(common_mul_mat_format_config_value(match.root.weight_format))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.has_bias", match.has_bias ? "1" : "0"); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.has_residual", match.has_residual ? "1" : "0"); + if (match.has_rmsnorm) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_postops.rms_epsilon", + common_mul_mat_id_to_config_value(match.epsilon)); + } + + dispatch.bindings.push_back({ match.root.input->id, 0, match.root.input->byte_count }); + dispatch.bindings.push_back({ routing_bundle.expert_table, 0, routing_bundle.expert_table_byte_count }); + dispatch.bindings.push_back({ routing_bundle.partition_table, 0, routing_bundle.partition_table_byte_count }); + dispatch.bindings.push_back({ match.root.weight->id, 0, match.root.weight->byte_count }); + if (match.has_bias) { + dispatch.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + } else { + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + } + if (match.has_residual) { + dispatch.bindings.push_back({ match.residual_input->id, 0, match.residual_input->byte_count }); + } else { + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + } + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + if (match.has_rmsnorm) { + 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_counters, 0, static_cast(completion_counter_count) * sizeof(int32_t) }); + } + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes)) { + return false; + } + for (const GraphNode * add_node : match.add_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, add_node, + dispatch_match.covered_nodes)) { + return false; + } + } + if (match.has_rmsnorm && (!append_covered_node_index_once(context.graph, context.covered_nodes, match.rms_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.mul_node, + dispatch_match.covered_nodes))) { + return false; + } + + if (match.has_rmsnorm) { + dispatch_match.completion_counter_requests.push_back({ completion_counters, + "common.mul_mat_id_postops.completion_counters", + static_cast(completion_counter_count) }); + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +void register_mul_mat_id_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.mul_mat_id_postops.f32_f32_wmma", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::Fused, + 180, + DispatchSource::Common, + match_mul_mat_id_postops_dispatch, + }); + registry.add({ + "common.mul_mat_id.f32_f32_wmma", + GGML_OP_MUL_MAT_ID, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Common, + match_mul_mat_id_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.h new file mode 100644 index 000000000000..931158c11c48 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-id.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_mul_mat_id_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..cc3848c42e10 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h @@ -0,0 +1,196 @@ +#pragma once + +#include "ggml.h" + +#include + +namespace ggml::hrx { + +enum class CommonMulMatWeightFormat { + Q1_0, + Q3K, + Q4K, + Q4KRow64, + Q5K, + Q6K, + Q6KRow64, + Q4_0, + Q4_1, + Q5_0, + Q5_1, + IQ1_S, + IQ1_M, + IQ2_XXS, + IQ2_XS, + IQ3_XXS, + IQ2_S, + IQ3_S, + IQ4_NL, + IQ4_XS, + Q8_0, + Q8_1, + F16, + BF16, + F32, +}; + +inline bool common_mul_mat_format_for_type(ggml_type type, CommonMulMatWeightFormat & format) { + switch (type) { + case GGML_TYPE_Q1_0: + format = CommonMulMatWeightFormat::Q1_0; + return true; + case GGML_TYPE_Q3_K: + format = CommonMulMatWeightFormat::Q3K; + return true; + case GGML_TYPE_Q4_K: + format = CommonMulMatWeightFormat::Q4K; + return true; + case GGML_TYPE_Q5_K: + format = CommonMulMatWeightFormat::Q5K; + return true; + case GGML_TYPE_Q6_K: + format = CommonMulMatWeightFormat::Q6K; + return true; + case GGML_TYPE_Q4_0: + format = CommonMulMatWeightFormat::Q4_0; + return true; + case GGML_TYPE_Q4_1: + format = CommonMulMatWeightFormat::Q4_1; + return true; + case GGML_TYPE_Q5_0: + format = CommonMulMatWeightFormat::Q5_0; + return true; + 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; + case GGML_TYPE_IQ3_S: + format = CommonMulMatWeightFormat::IQ3_S; + return true; + case GGML_TYPE_IQ4_NL: + format = CommonMulMatWeightFormat::IQ4_NL; + return true; + case GGML_TYPE_IQ4_XS: + format = CommonMulMatWeightFormat::IQ4_XS; + return true; + case GGML_TYPE_Q8_0: + format = CommonMulMatWeightFormat::Q8_0; + return true; + case GGML_TYPE_Q8_1: + format = CommonMulMatWeightFormat::Q8_1; + return true; + case GGML_TYPE_F16: + format = CommonMulMatWeightFormat::F16; + return true; + case GGML_TYPE_BF16: + format = CommonMulMatWeightFormat::BF16; + return true; + case GGML_TYPE_F32: + format = CommonMulMatWeightFormat::F32; + return true; + default: + return false; + } +} + +inline int64_t common_mul_mat_format_config_value(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::Q1_0: + return 10; + case CommonMulMatWeightFormat::Q3K: + return 11; + case CommonMulMatWeightFormat::Q4K: + return 4; + case CommonMulMatWeightFormat::Q4KRow64: + return 44; + case CommonMulMatWeightFormat::Q5K: + return 5; + case CommonMulMatWeightFormat::Q6K: + return 6; + case CommonMulMatWeightFormat::Q6KRow64: + return 46; + case CommonMulMatWeightFormat::Q4_0: + return 40; + case CommonMulMatWeightFormat::Q4_1: + return 41; + case CommonMulMatWeightFormat::Q5_0: + 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: + return 21; + case CommonMulMatWeightFormat::IQ4_NL: + return 20; + case CommonMulMatWeightFormat::IQ4_XS: + return 23; + case CommonMulMatWeightFormat::Q8_0: + return 80; + case CommonMulMatWeightFormat::Q8_1: + return 81; + case CommonMulMatWeightFormat::F16: + return 16; + case CommonMulMatWeightFormat::BF16: + return 30; + case CommonMulMatWeightFormat::F32: + return 32; + } + return 0; +} + +inline bool common_mul_mat_dense_float_format(CommonMulMatWeightFormat format) { + return format == CommonMulMatWeightFormat::F16 || format == CommonMulMatWeightFormat::BF16 || + format == CommonMulMatWeightFormat::F32; +} + +inline bool common_mul_mat_supported_dense_input_size(CommonMulMatWeightFormat format, int64_t input_size) { + switch (format) { + case CommonMulMatWeightFormat::Q1_0: + return input_size >= 256 && input_size <= 32768 && input_size % 128 == 0; + case CommonMulMatWeightFormat::Q4_0: + case CommonMulMatWeightFormat::Q4_1: + case CommonMulMatWeightFormat::Q5_0: + case CommonMulMatWeightFormat::Q5_1: + case CommonMulMatWeightFormat::IQ4_NL: + case CommonMulMatWeightFormat::Q8_0: + case CommonMulMatWeightFormat::Q8_1: + return input_size >= 256 && input_size <= 32768 && input_size % 32 == 0; + default: + return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; + } +} + +inline bool common_mul_mat_alternate_format_for_type(ggml_type type, CommonMulMatWeightFormat & format) { + switch (type) { + case GGML_TYPE_Q8_1: + case GGML_TYPE_F16: + case GGML_TYPE_F32: + return common_mul_mat_format_for_type(type, format); + default: + return false; + } +} + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..047caa6acec3 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp @@ -0,0 +1,2338 @@ +#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 +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kMulMatF32F32WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_wmma"); +static constexpr KernelCatalogRef kMulMatTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_tiled_input_f32_publish_f32"); +static constexpr KernelCatalogRef 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 = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_publish_f32"); +static constexpr KernelCatalogRef kMulMatSkinnyBiasAddF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"); +static constexpr KernelCatalogRef kMulMatF32F32DecodeWave64Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_f32_f32_decode_wave64"); +static constexpr KernelCatalogRef kMulMatAddF32F32DecodeWave64Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_add_f32_f32_decode_wave64"); +static constexpr KernelCatalogRef 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 = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_wmma"); +static constexpr KernelCatalogRef kQuantizeF32SymmetricI4K32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_quantize_f32_symmetric_i4_k32"); +static constexpr KernelCatalogRef kMulMatSymmetricI4LowRowWmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_wmma"); +static constexpr KernelCatalogRef kMulMatSymmetricI4LowRowSplitK2WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma"); +static constexpr KernelCatalogRef kMulMatSymmetricI4LowRowSplitK2DirectDotKernels[] = { + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c1"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c2"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c3"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c4"), + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c5"), +}; +static constexpr KernelCatalogRef kQuantizeF32SymmetricI8K256Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_quantize_f32_symmetric_i8_k256"); +static constexpr KernelCatalogRef kMulMatQ5KSymmetricI8WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q5_k_symmetric_i8_wmma"); +static constexpr KernelCatalogRef kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256"); +static constexpr KernelCatalogRef kMulMatQ4KF16WmmaPrefillWave32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); +static constexpr KernelCatalogRef kMulMatVectorQ6F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q6_f32_f32"); +static constexpr int64_t kMulMatVectorDenseMaxOutputSize = 524288; +static constexpr int64_t kMulMatVectorPackedMaxOutputSize = 262144; +static constexpr int64_t kMulMatVectorSelectedOutputsMaxOutputSize = 262144; +static constexpr KernelCatalogRef kMulMatQ6KF32WmmaPrefillWave32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_f32_wmma_prefill_wave32"); +static constexpr KernelCatalogRef kMulMatQ6KF16WmmaPrefillWave32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_f16_wmma_prefill_wave32"); +static constexpr KernelCatalogRef kSelectSymmetricI4K32GroupsKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_select_symmetric_i4_k32_groups"); +static constexpr KernelCatalogRef kMulMatQ6KSymmetricI2ScanToken1Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_q6_k_symmetric_i2_scan_token1"); +static constexpr KernelCatalogRef kTopK8F32PartitionsRegisterKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_top_k8_f32_partitions_register"); +static constexpr KernelCatalogRef kTopK128F32ReduceGatherRegisterKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_top_k128_f32_reduce_gather_register"); +static constexpr KernelCatalogRef kFillNegativeF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_fill_negative_f32"); +static constexpr KernelCatalogRef kMulMatVectorQ6SelectedOutputsF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_vector_q6_selected_outputs_f32_f32"); +static constexpr int64_t kQwenVectorVocabularyCount = 151936; + +static bool deprecated_mul_mat_dispatch_disabled() { + return std::getenv("GGML_HRX_DISABLE_DEPRECATED_MUL_MAT_DISPATCH") != nullptr; +} + +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; + const Value * projection_output = nullptr; + const Value * bias = nullptr; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const GraphNode * bias_add_node = nullptr; + const GraphNode * residual_add_node = nullptr; + const GraphNode * layout_node = nullptr; + std::vector add_nodes; + KernelCatalogRef kernel = {}; + CommonMulMatWeightFormat weight_format = CommonMulMatWeightFormat::Q4K; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + bool has_bias = false; + bool has_residual = false; + bool requires_q8_activation = false; + + bool matched() const { + return input != nullptr && weight != nullptr && projection_output != nullptr && residual_output != nullptr && + kernel.id != kUncatalogedKernelId && (has_bias || has_residual); + } +}; + +struct DecodeMulMatAddMatch { + CommonMulMatMatch root; + const Value * residual_input = nullptr; + const Value * residual_output = nullptr; + const GraphNode * add_node = nullptr; + + bool matched() const { + return root.matched() && residual_input != nullptr && residual_output != nullptr && add_node != nullptr; + } +}; + +struct VectorPublishPlan { + CommonActivationPublicationDemand demand; + CommonActivationPublication publication; +}; + +static bool common_mul_mat_is_low_token_batch(int64_t token_count) { + return token_count >= 2 && token_count <= 5; +} + +static bool common_mul_mat_has_skinny_reduction_shape(int64_t input_size) { + return input_size % 256 == 0; +} + +static bool common_mul_mat_is_skinny_route(const CommonMulMatMatch & match) { + return match.matched() && common_mul_mat_is_low_token_batch(match.token_count) && + common_mul_mat_has_skinny_reduction_shape(match.input_size); +} + +// Skinny owns the aligned 2..5 token optimization. Tiled owns normal prefill +// and the low-token tail fallback shapes. +static bool common_mul_mat_is_tiled_route(const CommonMulMatMatch & match) { + return match.matched() && match.token_count > 1 && !common_mul_mat_is_skinny_route(match); +} + +static bool common_mul_mat_postops_is_skinny_route(const MulMatPostOpsMatch & match) { + return match.matched() && common_mul_mat_is_low_token_batch(match.token_count) && + common_mul_mat_has_skinny_reduction_shape(match.input_size); +} + +static bool common_mul_mat_postops_is_tiled_route(const MulMatPostOpsMatch & match) { + return match.matched() && match.layout_node == nullptr && !match.requires_q8_activation && match.token_count > 1 && + !common_mul_mat_postops_is_skinny_route(match); +} + +static bool common_mul_mat_uses_factored_tiled_kernel(KernelCatalogRef kernel) { + return kernel.id == kMulMatTiledF32F32Kernel.id || kernel.id == kMulMatTiledF32F16AlternateKernel.id || + kernel.id == kMulMatTiledBiasAddF32F32Kernel.id; +} + +static bool is_bias_shape(const Value & value, int64_t output_size) { + if (value.kind != ValueKind::External || value.type != GGML_TYPE_F32 || !value.contiguous || + value.ne[0] != output_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] != 1) { + return false; + } + } + return true; +} + +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()); +} + +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 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; + } + 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; + } + ++compatible_consumer_count; + total_output_size += candidate_match.output_size; + } + 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) { + if (!graph.has_index()) { + return false; + } + for (const GraphNode * layout : graph.index().consumers(value.id)) { + if (layout == nullptr || (layout->op != GGML_OP_RESHAPE && layout->op != GGML_OP_VIEW)) { + continue; + } + const GraphNode * rms = common_find_single_consumer_with_op(graph, layout->output, GGML_OP_RMS_NORM); + const GraphNode * mul = + rms != nullptr ? common_find_single_consumer_with_op(graph, rms->output, GGML_OP_MUL) : nullptr; + if (mul != nullptr && common_find_single_consumer_with_op(graph, mul->output, GGML_OP_ROPE) != nullptr) { + return true; + } + } + return false; +} + +static bool has_reshaped_unary_mul_consumer(const Graph & graph, const Value & value) { + const GraphNode * reshape = common_find_single_consumer_with_op(graph, value.id, GGML_OP_RESHAPE); + const GraphNode * unary = + reshape != nullptr ? common_find_single_consumer_with_op(graph, reshape->output, GGML_OP_UNARY) : nullptr; + return unary != nullptr && common_find_single_consumer_with_op(graph, unary->output, GGML_OP_MUL) != nullptr; +} + +static bool has_reshaped_transpose_concat_consumer(const Graph & graph, const Value & value) { + const GraphNode * reshape = common_find_single_consumer_with_op(graph, value.id, GGML_OP_RESHAPE); + const GraphNode * transpose = + reshape != nullptr ? common_find_single_consumer_with_op(graph, reshape->output, GGML_OP_TRANSPOSE) : nullptr; + return transpose != nullptr && + common_find_single_consumer_with_op(graph, transpose->output, GGML_OP_CONCAT) != nullptr; +} + +static const Value * match_scaled_normalized_branch(const Graph & graph, ValueId branch_id, int64_t hidden_size) { + const Value * branch = common_graph_value(graph, branch_id); + const GraphNode * mul = graph.index().producer(branch_id); + if (branch == nullptr || mul == nullptr || !common_binary_node_is_mul(*mul) || branch->type != GGML_TYPE_F32 || + !common_is_2d(*branch) || branch->ne[0] != hidden_size || branch->ne[1] != 1) { + return nullptr; + } + + const GraphNode * rms = nullptr; + const Value * scale = nullptr; + for (ValueId input_id : mul->inputs) { + const GraphNode * producer = graph.index().producer(input_id); + if (producer != nullptr && producer->op == GGML_OP_RMS_NORM) { + if (rms != nullptr) { + return nullptr; + } + rms = producer; + } else { + scale = common_graph_value(graph, input_id); + } + } + if (rms == nullptr || rms->inputs.size() != 1 || scale == nullptr || scale->type != GGML_TYPE_F32 || + !common_is_2d(*scale) || scale->ne[0] != hidden_size || scale->ne[1] != 1) { + return nullptr; + } + + const Value * normalized = common_graph_value(graph, rms->output); + const Value * source = common_graph_value(graph, rms->inputs[0]); + if (normalized == nullptr || source == nullptr || normalized->type != GGML_TYPE_F32 || + source->type != GGML_TYPE_F32 || !common_is_2d(*normalized) || !common_is_2d(*source) || + normalized->ne[0] != hidden_size || normalized->ne[1] != 1 || source->ne[0] != hidden_size || + source->ne[1] != 1) { + return nullptr; + } + return source; +} + +static bool matches_embedded_external_concat_projection(const Graph & graph, + const GraphNode & projection, + int64_t hidden_size) { + if (projection.op != GGML_OP_MUL_MAT || projection.inputs.size() != 2) { + return false; + } + + const Value * weight = common_graph_value(graph, projection.inputs[0]); + const Value * input = common_graph_value(graph, projection.inputs[1]); + const Value * output = common_graph_value(graph, projection.output); + if (weight == nullptr || input == nullptr || output == nullptr || !common_is_2d(*weight) || !common_is_2d(*input) || + !common_is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || + input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || weight->ne[0] != 2 * hidden_size || + weight->ne[1] != hidden_size || input->ne[0] != 2 * hidden_size || input->ne[1] != 1 || + output->ne[0] != hidden_size || output->ne[1] != 1) { + return false; + } + + const GraphNode * concat = graph.index().producer(input->id); + if (concat == nullptr || concat->op != GGML_OP_CONCAT || concat->inputs.size() != 2) { + return false; + } + const Value * first_source = match_scaled_normalized_branch(graph, concat->inputs[0], hidden_size); + const Value * second_source = match_scaled_normalized_branch(graph, concat->inputs[1], hidden_size); + if (first_source == nullptr || second_source == nullptr) { + return false; + } + + const GraphNode * first_producer = graph.index().producer(first_source->id); + const GraphNode * second_producer = graph.index().producer(second_source->id); + const bool first_is_embedding = + first_producer != nullptr && first_producer->op == GGML_OP_GET_ROWS && first_producer->inputs.size() == 2; + const bool second_is_embedding = + second_producer != nullptr && second_producer->op == GGML_OP_GET_ROWS && second_producer->inputs.size() == 2; + const bool first_is_external = first_producer == nullptr; + const bool second_is_external = second_producer == nullptr; + return (first_is_embedding && second_is_external) || (second_is_embedding && first_is_external); +} + +static bool has_embedded_external_concat_projection_ancestor(const Graph & graph, + ValueId endpoint_input, + int64_t hidden_size) { + if (!graph.has_index() || endpoint_input.value < 0 || hidden_size <= 0) { + return false; + } + + std::vector visited(graph.values().size(), false); + std::vector pending = { endpoint_input }; + while (!pending.empty()) { + const ValueId current = pending.back(); + pending.pop_back(); + if (current.value < 0 || static_cast(current.value) >= visited.size() || visited[current.value]) { + continue; + } + visited[current.value] = true; + const GraphNode * producer = graph.index().producer(current); + if (producer == nullptr) { + continue; + } + if (matches_embedded_external_concat_projection(graph, *producer, hidden_size)) { + return true; + } + pending.insert(pending.end(), producer->inputs.begin(), producer->inputs.end()); + } + return false; +} + +static bool is_qwen_endpoint_rmsnorm_projection(const Graph & graph, const CommonMulMatMatch & match) { + if (!graph.has_index() || match.weight == nullptr || match.input == nullptr || + match.weight->type != GGML_TYPE_Q6_K || match.token_count != 1 || + match.output_size != kQwenVectorVocabularyCount) { + return false; + } + + const GraphNode * mul = graph.index().producer(match.input->id); + if (mul == nullptr || mul->op != GGML_OP_MUL || mul->inputs.size() != 2) { + return false; + } + for (ValueId input : mul->inputs) { + const GraphNode * producer = graph.index().producer(input); + if (producer != nullptr && producer->op == GGML_OP_RMS_NORM) { + return true; + } + } + return false; +} + +static bool match_symmetric_i4_low_row_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + 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() || + !common_symmetric_i4_lowrow_mul_mat_eligible(context.graph, *context.root_node, *match.input)) { + return false; + } + const CommonSymmetricI4ActivationLayout activation_layout = + common_symmetric_i4_activation_layout(match.input_size, match.token_count); + const CommandPlanAlternateValue * alternate = find_alternate_value(context.graph, context.plan, match.input->id, + GGML_TYPE_COUNT, activation_layout.total_bytes); + if (alternate != nullptr && alternate->name != kCommonSymmetricI4K32ActivationAlternateName) { + alternate = nullptr; + } + + const ValueId activation = alternate != nullptr ? alternate->alternate_value : context.next_plan_value; + const int32_t next_value = context.next_plan_value.value + (alternate == nullptr ? 1 : 0); + + if (alternate == nullptr) { + dispatch_match.transients.push_back( + { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeF32SymmetricI4K32Kernel); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.input_size", + common_to_config_value(match.input_size)); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.token_count", + common_to_config_value(match.token_count)); + quantize.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + quantize.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + quantize.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + quantize.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + dispatch_match.dispatches.push_back(std::move(quantize)); + + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.input->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, + kCommonSymmetricI4K32ActivationAlternateName }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + } + + const size_t materialized_weight_bytes = + common_symmetric_i4_shared4_weight_byte_count(match.input_size, match.output_size); + const bool split_k = match.input_size % 128 == 0 && materialized_weight_bytes >= size_t{ 16 } * 1024 * 1024 && + match.output_size <= 2 * match.input_size; + + const ValueId partial = split_k ? ValueId(next_value) : ValueId{}; + const ValueId completion_counters = split_k ? ValueId(next_value + 1) : ValueId{}; + const uint32_t completion_counter_count = + static_cast(common_ceil_div(match.output_size, 16) * common_ceil_div(match.token_count, 16)); + if (split_k) { + dispatch_match.transients.push_back( + { partial, "common.mul_mat.symmetric_i4_lowrow.split_k_partial", match.output->byte_count, 256 }); + dispatch_match.completion_counter_requests.push_back({ + completion_counters, + "common.mul_mat.symmetric_i4_lowrow.split_k_completion_counters", + completion_counter_count, + }); + } + + const bool use_split_direct_dot = split_k && match.token_count <= 5 && match.input_size % 256 == 0; + + Dispatch contraction; + contraction.kernel = make_kernel_specialization( + use_split_direct_dot ? kMulMatSymmetricI4LowRowSplitK2DirectDotKernels[match.token_count - 1] : + split_k ? kMulMatSymmetricI4LowRowSplitK2WmmaKernel : + kMulMatSymmetricI4LowRowWmmaKernel); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.input_size", + common_to_config_value(match.input_size)); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.output_size", + common_to_config_value(match.output_size)); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.token_count", + common_to_config_value(match.token_count)); + contraction.kernel.compile_parameters.emplace( + "ggml.mul_mat.symmetric_i4.lowrow.row_group_size", + common_to_config_value( + static_cast(common_symmetric_shared4_row_group_size(match.input_size, match.output_size, 4)))); + contraction.bindings.push_back( + match.weight->type == GGML_TYPE_Q5_K && match.token_count == 1 ? + common_symmetric_i4_shared4_multistart_weight_binding(*match.weight, match.input_size, match.output_size) : + common_symmetric_i4_shared4_weight_binding(*match.weight, match.input_size, match.output_size)); + contraction.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + contraction.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + contraction.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + contraction.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + contraction.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + if (split_k) { + contraction.bindings.push_back({ partial, 0, match.output->byte_count }); + contraction.bindings.push_back({ completion_counters, 0, completion_counter_count * sizeof(int32_t) }); + } + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(contraction)); + return dispatch_match.status.success(); +} + +static bool match_q6_vector_selected_outputs_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + constexpr int64_t selected_group_count = 96; + constexpr int64_t selected_output_count = 128; + constexpr int64_t selected_output_chunk_size = 64; + constexpr size_t partition_entry_count = 1024; + + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KSymmetricI2ScanToken1Kernel, true); + if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || + match.token_count != 1 || match.weight->alias_source.value >= 0 || match.input_size % 256 != 0 || + match.input_size > 8192 || match.output_size < 65536 || match.output_size % 64 != 0 || + match.output_size > kMulMatVectorSelectedOutputsMaxOutputSize || match.output_size < 16 * match.input_size || + !context.graph.index().consumers(match.output->id).empty() || + !has_embedded_external_concat_projection_ancestor(context.graph, match.input->id, match.input_size)) { + return false; + } + + const CommonSymmetricI4ActivationLayout activation_layout = + common_symmetric_i4_activation_layout(match.input_size, match.token_count); + const CommandPlanAlternateValue * alternate = find_alternate_value(context.graph, context.plan, match.input->id, + GGML_TYPE_COUNT, activation_layout.total_bytes); + if (alternate != nullptr && alternate->name != kCommonSymmetricI4K32ActivationAlternateName) { + alternate = nullptr; + } + + const ValueId activation = alternate != nullptr ? alternate->alternate_value : context.next_plan_value; + const int32_t transient_base = context.next_plan_value.value + (alternate == nullptr ? 1 : 0); + const ValueId selected_groups(transient_base); + const ValueId partial_values(transient_base + 1); + const ValueId partial_ids(transient_base + 2); + const ValueId selected_outputs(transient_base + 3); + const ValueId selected_values(transient_base + 4); + + if (alternate == nullptr) { + dispatch_match.transients.push_back( + { activation, kCommonSymmetricI4K32ActivationAlternateName, activation_layout.total_bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeF32SymmetricI4K32Kernel); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.input_size", + common_to_config_value(match.input_size)); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k32.token_count", + common_to_config_value(match.token_count)); + quantize.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + quantize.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + quantize.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + quantize.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + dispatch_match.dispatches.push_back(std::move(quantize)); + + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { match.input->id, activation, GGML_TYPE_COUNT, activation_layout.total_bytes, + kCommonSymmetricI4K32ActivationAlternateName }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + } + + dispatch_match.transients.push_back({ selected_groups, "common.mul_mat.vector_q6_selected_outputs.groups", + static_cast(selected_group_count) * sizeof(int32_t), 256 }); + dispatch_match.transients.push_back({ partial_values, "common.mul_mat.vector_q6_selected_outputs.partial_values", + partition_entry_count * sizeof(float), 256 }); + dispatch_match.transients.push_back({ partial_ids, "common.mul_mat.vector_q6_selected_outputs.partial_ids", + partition_entry_count * sizeof(int32_t), 256 }); + dispatch_match.transients.push_back({ selected_outputs, "common.mul_mat.vector_q6_selected_outputs.indices", + selected_output_count * sizeof(int32_t), 256 }); + dispatch_match.transients.push_back({ selected_values, "common.mul_mat.vector_q6_selected_outputs.values", + selected_output_count * sizeof(float), 256 }); + + Dispatch select_groups; + select_groups.kernel = make_kernel_specialization(kSelectSymmetricI4K32GroupsKernel); + select_groups.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_shortlist.input_size", + common_to_config_value(match.input_size)); + select_groups.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_shortlist.selected_group_count", + common_to_config_value(selected_group_count)); + select_groups.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + select_groups.bindings.push_back( + { selected_groups, 0, static_cast(selected_group_count) * sizeof(int32_t) }); + dispatch_match.dispatches.push_back(std::move(select_groups)); + + const size_t symmetric_i2_bytes = + common_symmetric_shared4_weight_byte_count(match.input_size, match.output_size, 2); + if (symmetric_i2_bytes > std::numeric_limits::max() - match.weight->byte_count) { + return false; + } + const size_t materialized_weight_bytes = symmetric_i2_bytes + match.weight->byte_count; + Dispatch scan; + scan.kernel = make_kernel_specialization(kMulMatQ6KSymmetricI2ScanToken1Kernel); + scan.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_shortlist.input_size", + common_to_config_value(match.input_size)); + scan.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_shortlist.output_size", + common_to_config_value(match.output_size)); + scan.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_shortlist.selected_group_count", + common_to_config_value(selected_group_count)); + scan.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, + kQ6KSymmetricI2PackedK256Row64ScaleRowLayout, match.weight->type, match.input_size, + match.output_size, match.weight->byte_count }); + scan.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + scan.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + scan.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + scan.bindings.push_back({ selected_groups, 0, static_cast(selected_group_count) * sizeof(int32_t) }); + dispatch_match.dispatches.push_back(std::move(scan)); + + Dispatch partition_top_k; + partition_top_k.kernel = make_kernel_specialization(kTopK8F32PartitionsRegisterKernel); + partition_top_k.kernel.integer_parameters.emplace("element_count", match.output_size); + partition_top_k.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + partition_top_k.bindings.push_back({ partial_values, 0, partition_entry_count * sizeof(float) }); + partition_top_k.bindings.push_back({ partial_ids, 0, partition_entry_count * sizeof(int32_t) }); + dispatch_match.dispatches.push_back(std::move(partition_top_k)); + + Dispatch reduce_top_k; + reduce_top_k.kernel = make_kernel_specialization(kTopK128F32ReduceGatherRegisterKernel); + reduce_top_k.kernel.integer_parameters.emplace("element_count", match.output_size); + reduce_top_k.bindings.push_back({ partial_values, 0, partition_entry_count * sizeof(float) }); + reduce_top_k.bindings.push_back({ partial_ids, 0, partition_entry_count * sizeof(int32_t) }); + reduce_top_k.bindings.push_back({ selected_outputs, 0, selected_output_count * sizeof(int32_t) }); + reduce_top_k.bindings.push_back({ selected_values, 0, selected_output_count * sizeof(float) }); + dispatch_match.dispatches.push_back(std::move(reduce_top_k)); + + Dispatch fill; + fill.kernel = make_kernel_specialization(kFillNegativeF32Kernel); + fill.kernel.integer_parameters.emplace("element_count", match.output_size); + fill.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch_match.dispatches.push_back(std::move(fill)); + + for (int64_t selected_output_offset = 0; selected_output_offset < selected_output_count; + selected_output_offset += selected_output_chunk_size) { + Dispatch selected_output_pass; + selected_output_pass.kernel = make_kernel_specialization(kMulMatVectorQ6SelectedOutputsF32Kernel); + selected_output_pass.kernel.integer_parameters.emplace("token_count", match.token_count); + selected_output_pass.kernel.integer_parameters.emplace("selected_count", selected_output_chunk_size); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.input_size", + common_to_config_value(match.input_size)); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.output_size", + common_to_config_value(match.output_size)); + selected_output_pass.kernel.compile_parameters.emplace("ggml.matmul.vector.output_accumulation", "0"); + selected_output_pass.kernel.compile_parameters.emplace( + "ggml.matmul.vector.weight_offset", common_to_config_value(static_cast(symmetric_i2_bytes))); + selected_output_pass.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + selected_output_pass.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, + kQ6KSymmetricI2PackedK256Row64ScaleRowLayout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + selected_output_pass.bindings.push_back({ selected_outputs, + static_cast(selected_output_offset) * sizeof(int32_t), + static_cast(selected_output_chunk_size) * sizeof(int32_t) }); + selected_output_pass.bindings.push_back({ selected_values, + static_cast(selected_output_offset) * sizeof(float), + static_cast(selected_output_chunk_size) * sizeof(float) }); + selected_output_pass.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch_match.dispatches.push_back(std::move(selected_output_pass)); + } + + dispatch_match.covered_nodes.push_back(context.root_index); + return dispatch_match.status.success(); +} + +enum class Q6VectorStorage : int64_t { + Packed = 0, + I8 = 1, +}; + +static bool build_q6_vector_dispatch(const CommonMulMatMatch & match, + size_t root_index, + Q6VectorStorage storage, + size_t weight_bytes, + DispatchMatch & dispatch_match) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatVectorQ6F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", 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.q6_storage", + common_to_config_value(static_cast(storage))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + if (storage == Q6VectorStorage::I8) { + dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KI8K32Row64Layout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, weight_bytes, kQ6KPackedK256Row64ScaleRowLayout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool build_q6_vector_dispatch(const CommonMulMatMatch & match, + size_t root_index, + DispatchMatch & dispatch_match) { + return build_q6_vector_dispatch(match, root_index, Q6VectorStorage::Packed, match.weight->byte_count, + dispatch_match); +} + +static bool match_q6_vector_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatVectorQ6F32Kernel, true); + if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || + match.weight->alias_source.value >= 0 || match.token_count != 1 || match.input_size % 256 != 0 || + match.output_size % 64 != 0 || match.output_size > kMulMatVectorPackedMaxOutputSize || + !context.graph.index().consumers(match.output->id).empty()) { + return false; + } + if (is_qwen_endpoint_rmsnorm_projection(context.graph, match)) { + return false; + } + if (find_alternate_value( + context.graph, context.plan, match.input->id, GGML_TYPE_Q8_1, + static_cast(match.token_count) * ggml_row_size(GGML_TYPE_Q8_1, match.input_size)) != nullptr) { + return false; + } + + return build_q6_vector_dispatch(match, context.root_index, dispatch_match); +} + +static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, + ValueId input, + size_t input_bytes, + KernelCatalogRef kernel, + size_t root_index, + DispatchMatch & dispatch_match) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_accumulation", "0"); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.token_capacity", + common_to_config_value(match.token_count)); + dispatch.bindings.push_back({ input, 0, input_bytes }); + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KF32WmmaPrefillWave32Kernel, false); + if (!match.matched() || match.weight->type != GGML_TYPE_Q6_K || match.weight->alias_source.value >= 0 || + match.token_count < 128 || match.token_count > 2048 || match.token_count % 128 != 0 || + match.input_size % 256 != 0 || match.output_size % 64 != 0) { + return false; + } + + const bool packed_input = + common_mul_mat_uses_k16_major_f16(match.weight_format, match.input_size, match.output_size, match.token_count); + DispatchBinding activation; + const bool prepared = packed_input ? + common_prepare_k16_major_f16_input(context, *match.input, match.input_size, + match.token_count, dispatch_match, activation) : + common_prepare_f16_input(context, *match.input, match.input_size, match.token_count, + dispatch_match, activation); + if (!prepared) { + return false; + } + return build_q6_k_prefill_wave32_dispatch(match, activation.value, activation.length, + kMulMatQ6KF16WmmaPrefillWave32Kernel, context.root_index, dispatch_match); +} + +static bool match_symmetric_i4_prefill_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; + + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSymmetricI4WmmaKernel, false); + if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q5_K || + match.weight->alias_source.value >= 0 || match.token_count < 128 || match.token_count % 128 != 0 || + match.input_size % 64 != 0 || match.output_size < match.input_size || match.output_size % 128 != 0) { + return false; + } + + const GraphNode * input_producer = context.graph.index().producer(match.input->id); + if (input_producer != nullptr && input_producer->op == GGML_OP_GLU) { + return false; + } + if (!has_normalized_rope_consumer(context.graph, *match.output) && + !has_reshaped_unary_mul_consumer(context.graph, *match.output) && + !has_reshaped_transpose_concat_consumer(context.graph, *match.output)) { + return false; + } + + const CommonSymmetricI4ActivationLayout activation_layout = + common_symmetric_i4_activation_layout(match.input_size, match.token_count); + const ValueId activation = context.next_plan_value; + dispatch_match.transients.push_back( + { activation, "common.mul_mat.symmetric_i4_k64.activation", activation_layout.total_bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeF32SymmetricI4K64PlaneKernel); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k64.input_size", + common_to_config_value(match.input_size)); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i4_k64.token_count", + common_to_config_value(match.token_count)); + quantize.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + quantize.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + quantize.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + quantize.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + + const size_t materialized_weight_bytes = + static_cast(match.output_size) * static_cast(match.input_size / 256) * size_t{ 144 }; + Dispatch contraction; + contraction.kernel = make_kernel_specialization(kMulMatSymmetricI4WmmaKernel); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.input_size", + common_to_config_value(match.input_size)); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.output_size", + common_to_config_value(match.output_size)); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.token_count", + common_to_config_value(match.token_count)); + contraction.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, kSymmetricI4K64Row64Layout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + contraction.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + contraction.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + contraction.bindings.push_back({ activation, 0, activation_layout.payload_bytes }); + contraction.bindings.push_back({ activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + contraction.bindings.push_back({ activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(quantize)); + dispatch_match.dispatches.push_back(std::move(contraction)); + return true; +} + +static bool match_q5_k_symmetric_i8_prefill_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; + + const CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ5KSymmetricI8WmmaKernel, false); + if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q5_K || + match.weight->alias_source.value >= 0 || match.token_count < 256 || match.token_count % 256 != 0 || + match.input_size % 256 != 0 || match.output_size % 64 != 0 || match.output_size > match.input_size) { + return false; + } + + const GraphNode * input_producer = context.graph.index().producer(match.input->id); + if (input_producer != nullptr && input_producer->op == GGML_OP_GLU) { + return false; + } + + const size_t element_count = static_cast(match.token_count) * static_cast(match.input_size); + const size_t metadata_bytes = + static_cast(match.token_count) * static_cast(match.input_size / 256) * sizeof(int32_t); + const size_t quantized_bytes = element_count + metadata_bytes; + const ValueId quantized = context.next_plan_value; + dispatch_match.transients.push_back( + { quantized, "common.mul_mat.symmetric_i8_k256.activation", quantized_bytes, 256 }); + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeF32SymmetricI8K256Kernel); + quantize.kernel.integer_parameters.emplace("token_count", match.token_count); + quantize.kernel.integer_parameters.emplace("input_size_arg", match.input_size); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i8_k256.token_count", + common_to_config_value(match.token_count)); + quantize.kernel.compile_parameters.emplace("ggml.quantize_symmetric_i8_k256.input_size", + common_to_config_value(match.input_size)); + quantize.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + quantize.bindings.push_back({ quantized, 0, quantized_bytes }); + + const size_t materialized_weight_bytes = + static_cast(match.output_size) * static_cast(match.input_size / 256) * size_t{ 258 }; + Dispatch contraction; + contraction.kernel = make_kernel_specialization(kMulMatQ5KSymmetricI8WmmaKernel); + contraction.kernel.integer_parameters.emplace("token_count", match.token_count); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i8.input_size", + common_to_config_value(match.input_size)); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i8.output_size", + common_to_config_value(match.output_size)); + contraction.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i8.token_count", + common_to_config_value(match.token_count)); + contraction.bindings.push_back({ quantized, 0, quantized_bytes }); + contraction.bindings.push_back({ match.weight->id, 0, materialized_weight_bytes, kQ5KSymmetricI8K256Row64Layout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + contraction.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(quantize)); + dispatch_match.dispatches.push_back(std::move(contraction)); + return true; +} + +static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, + kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel, false); + if (!match.matched() || !context.graph.has_index() || + (match.weight->type != GGML_TYPE_Q4_K && match.weight->type != GGML_TYPE_Q5_K && + match.weight->type != GGML_TYPE_IQ4_XS) || + match.token_count < 256 || match.token_count > 2048 || match.token_count % 256 != 0 || + match.input_size % 256 != 0 || match.output_size % 64 != 0) { + return false; + } + + for (const GraphNode * consumer : context.graph.index().consumers(match.output->id)) { + if (match.weight->type != GGML_TYPE_Q4_K && consumer != nullptr && consumer->op == GGML_OP_GLU) { + return false; + } + } + + const bool use_f16 = match.weight->type == GGML_TYPE_Q4_K && match.weight->alias_source.value < 0 && + match.output_size >= match.input_size / 4; + DispatchBinding activation; + const bool packed_input = + use_f16 && + common_mul_mat_uses_k16_major_f16(match.weight_format, match.input_size, match.output_size, match.token_count, + context.plan.metadata.find_generated_resource( + match.input->id, GeneratedResourceRole::F16K16Major) != nullptr); + if (use_f16) { + const bool prepared = packed_input ? + common_prepare_k16_major_f16_input(context, *match.input, match.input_size, + match.token_count, dispatch_match, activation) : + common_prepare_f16_input(context, *match.input, match.input_size, match.token_count, + dispatch_match, activation); + if (!prepared) { + return false; + } + } else if (!common_prepare_q8_1_x4_input( + context, *match.input, match.input_size, match.token_count, dispatch_match, activation, + match.weight->type == GGML_TYPE_Q4_K ? CommonQ8ActivationPolicy::AllowStandaloneQuantize : + CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(use_f16 ? kMulMatQ4KF16WmmaPrefillWave32Kernel : + kMulMatQ5KIQ4XSQ8_1X4WmmaToken256Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace(use_f16 ? "ggml.mul_mat.input_size" : "ggml.mul_mat_q8_1_x4.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace( + use_f16 ? "ggml.mul_mat.output_size" : "ggml.mul_mat_q8_1_x4.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace( + use_f16 ? "ggml.workload.token_capacity" : "ggml.mul_mat_q8_1_x4.token_capacity", + common_to_config_value(match.token_count)); + const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && match.weight->alias_source.value < 0; + if (use_f16) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.f16_input_layout", packed_input ? "1" : "0"); + } else { + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_q8_1_x4.weight_format", + common_to_config_value(common_mul_mat_format_config_value(pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : + match.weight_format))); + } + dispatch.bindings.push_back(activation); + if (pack_q4) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, kQ4KPackedK256Row64Layout, + match.weight->type, match.input_size, match.output_size, + match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_q4_k_q8_1_x4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + if (context.root_node == nullptr || context.root_node->inputs.empty()) { + return false; + } + const Value * weight = common_graph_value(context.graph, context.root_node->inputs[0]); + return weight != nullptr && weight->type == GGML_TYPE_Q4_K && + match_packed_q8_1_x4_prefill_dispatch(context, dispatch_match); +} + +static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, + KernelCatalogRef root_kernel, + KernelCatalogRef bias_kernel, + KernelCatalogRef add_kernel, + KernelCatalogRef bias_add_kernel, + bool allow_batch1_bias_residual = false) { + MulMatPostOpsMatch match; + CommonMulMatMatch root = common_match_mul_mat_any_format(context.graph, context.root_node, root_kernel, false); + if (!root.matched()) { + root = common_match_mul_mat_any_format(context.graph, context.root_node, root_kernel, true); + } + if (!root.matched() || !context.graph.has_index()) { + return match; + } + + const bool packed_row_contraction = + root.token_count <= 5 && root.weight->alias_source.value < 0 && + (root.weight_format == CommonMulMatWeightFormat::Q4K || root.weight_format == CommonMulMatWeightFormat::Q6K) && + root.input_size >= 4096 && root.output_size >= 4096 && root.output_size <= root.input_size && + root.output_size % 64 == 0; + const bool batch1_requires_q8_activation = root.token_count == 1 && !packed_row_contraction; + + const Value * current = root.output; + if (packed_row_contraction) { + const GraphNode * reshape = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_RESHAPE); + const Value * reshaped = reshape != nullptr ? common_graph_value(context.graph, reshape->output) : nullptr; + if (reshaped != nullptr && is_layout_alias_node(context.graph, *reshape) && reshaped->contiguous && + same_full_value_range(*current, *reshaped) && common_same_shape(*current, *reshaped)) { + match.layout_node = reshape; + current = reshaped; + } + } + const int add_limit = + match.layout_node != nullptr || (root.token_count == 1 && !allow_batch1_bias_residual) ? 1 : 2; + for (int add_index = 0; add_index < add_limit; ++add_index) { + const GraphNode * add_node = common_find_only_consumer_with_op(context.graph, current->id, GGML_OP_ADD); + if (add_node == nullptr || !common_binary_node_is_add(*add_node)) { + break; + } + + const bool current_is_lhs = add_node->inputs[0] == current->id; + const bool current_is_rhs = add_node->inputs[1] == current->id; + if (!current_is_lhs && !current_is_rhs) { + return {}; + } + + const ValueId other_id = current_is_lhs ? add_node->inputs[1] : add_node->inputs[0]; + const Value * other = common_graph_value(context.graph, other_id); + const Value * output = common_graph_value(context.graph, add_node->output); + if (other == nullptr || output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || + !common_same_shape(*root.output, *output)) { + return {}; + } + + if (!match.has_bias && is_bias_shape(*other, root.output_size)) { + match.bias = other; + match.bias_add_node = add_node; + match.add_nodes.push_back(add_node); + match.has_bias = true; + current = output; + continue; + } + + if (!match.has_residual && other->type == GGML_TYPE_F32 && other->contiguous && + common_same_shape(*root.output, *other)) { + match.residual_input = other; + match.residual_output = output; + match.residual_add_node = add_node; + match.add_nodes.push_back(add_node); + match.has_residual = true; + current = output; + continue; + } + + break; + } + + if (!match.has_bias && !match.has_residual) { + return {}; + } + + match.residual_output = current; + + if (match.has_bias && match.has_residual) { + match.kernel = bias_add_kernel; + } else if (match.has_residual) { + match.kernel = add_kernel; + } else if (match.has_bias) { + match.kernel = bias_kernel; + } else { + return {}; + } + + match.input = root.input; + match.weight = root.weight; + match.projection_output = root.output; + match.weight_format = root.weight_format; + match.input_size = root.input_size; + match.output_size = root.output_size; + match.token_count = root.token_count; + match.requires_q8_activation = batch1_requires_q8_activation; + return match; +} + +static MulMatPostOpsMatch match_mul_mat_postops(const DispatchMatchContext & context, bool allow_batch1_bias_residual) { + return match_mul_mat_postops(context, kMulMatF32F32WmmaKernel, kMulMatVectorBiasAddF32F32Kernel, + kMulMatVectorBiasAddF32F32Kernel, kMulMatVectorBiasAddF32F32Kernel, + allow_batch1_bias_residual); +} + +static bool try_match_fused_unary(const DispatchMatchContext & context, CommonMulMatMatch & match) { + if (!match.matched()) { + return false; + } + + const std::vector & consumers = context.graph.index().consumers(context.root_node->output); + if (consumers.size() != 1 || consumers.front() == nullptr) { + return false; + } + + const GraphNode * unary = consumers.front(); + size_t unary_index = 0; + if (!context.graph.index().node_index(unary, unary_index) || unary_index >= context.covered_nodes.size() || + context.covered_nodes[unary_index] || unary->inputs.size() != 1) { + return false; + } + + const UnaryParams * params = op_params_as(unary->params); + if (params == nullptr || !unary_kind_supported(params->op)) { + return false; + } + + const Value * unary_input = common_graph_value(context.graph, unary->inputs[0]); + const Value * unary_output = common_graph_value(context.graph, unary->output); + if (unary_input == nullptr || unary_output == nullptr || unary_input->id != match.output->id || + unary_output->type != GGML_TYPE_F32 || !unary_output->contiguous || + !common_same_shape(*match.output, *unary_output)) { + return false; + } + + match.output = unary_output; + match.output_unary_op = params->op; + match.unary_node_index = unary_index; + match.has_fused_unary = true; + return true; +} + +} // namespace + +static bool build_mul_mat_dispatch( + const DispatchMatchContext & context, + const CommonMulMatMatch & match, + DispatchMatch & dispatch_match, + CommonQ8ActivationPolicy q8_policy = CommonQ8ActivationPolicy::ExistingAlternateOnly) { + 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( + 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); + 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"); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.output_unary_op", + std::to_string(unary_kind_config_value(match.output_unary_op))); + const bool pack_rows = match.token_count <= 5 && common_mul_mat_has_skinny_reduction_shape(match.input_size) && + match.output_size % 64 == 0 && match.weight->alias_source.value < 0; + const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; + const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; + bool pack_q4 = false; + bool pack_q6 = false; + DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; + const size_t f16_activation_bytes = static_cast(match.token_count * match.input_size) * sizeof(ggml_fp16_t); + const CommandPlanAlternateValue * f16_activation = + common_mul_mat_uses_factored_tiled_kernel(match.kernel) ? + find_alternate_value(context.graph, context.plan, match.input->id, GGML_TYPE_F16, f16_activation_bytes) : + nullptr; + if (f16_activation != nullptr) { + activation = { f16_activation->alternate_value, 0, f16_activation_bytes }; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); + } else if (can_pack_q4 || can_pack_q6) { + if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + activation, q8_policy)) { + pack_q4 = can_pack_q4; + pack_q6 = can_pack_q6; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", + std::to_string(GGML_TYPE_Q8_1)); + } else if (q8_policy == CommonQ8ActivationPolicy::AllowStandaloneQuantize) { + return false; + } + } + const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : + pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : + match.weight_format; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.weight_format", + common_to_config_value(common_mul_mat_format_config_value(format))); + dispatch.bindings.push_back(activation); + if (pack_q4 || pack_q6) { + const char * layout = pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, layout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + if (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) { + constexpr const char * name = "common.mul_mat.tiled.f16"; + 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(publication.binding()); + } + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.has_fused_unary) { + dispatch_match.covered_nodes.push_back(match.unary_node_index); + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + if (publish_q8_alternate) { + constexpr const char * name = "common.mul_mat.tiled.q8_1_x4"; + 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); + quantize.kernel.integer_parameters.emplace("token_count", match.token_count); + quantize.kernel.integer_parameters.emplace("input_size", match.output_size); + quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + common_to_config_value(match.token_count * match.output_size / 128)); + quantize.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + quantize.bindings.push_back(publication.binding()); + dispatch_match.dispatches.push_back(std::move(quantize)); + } + return true; +} + +static bool prepare_vector_activation_binding(const DispatchMatchContext & context, + const Value & input, + const Value & weight, + int64_t input_size, + int64_t output_size, + int64_t token_count, + CommonMulMatWeightFormat weight_format, + DispatchMatch & dispatch_match, + DispatchBinding & activation, + CommonMulMatWeightFormat & kernel_weight_format, + const char *& weight_layout, + bool & uses_q8_activation) { + activation = { input.id, 0, input.byte_count }; + kernel_weight_format = weight_format; + weight_layout = nullptr; + uses_q8_activation = false; + + if (token_count != 1 || input_size % 256 != 0 || weight.alias_source.value >= 0 || output_size % 64 != 0) { + return true; + } + + const bool can_pack_q4 = weight_format == CommonMulMatWeightFormat::Q4K; + const bool can_pack_q6 = weight_format == CommonMulMatWeightFormat::Q6K; + if (!can_pack_q4 && !can_pack_q6) { + return true; + } + + DispatchBinding q8_activation = activation; + if (!common_prepare_q8_1_x4_input(context, input, input_size, token_count, dispatch_match, q8_activation, + CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + return true; + } + + activation = q8_activation; + kernel_weight_format = can_pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; + weight_layout = can_pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + uses_q8_activation = true; + return true; +} + +static bool is_k256_packed_row_layout(CommonMulMatWeightFormat format) { + return format == CommonMulMatWeightFormat::Q4KRow64 || format == CommonMulMatWeightFormat::Q6KRow64; +} + +static int64_t vector_max_output_size(CommonMulMatWeightFormat format) { + switch (format) { + case CommonMulMatWeightFormat::Q4K: + case CommonMulMatWeightFormat::Q4KRow64: + case CommonMulMatWeightFormat::Q6K: + case CommonMulMatWeightFormat::Q6KRow64: + return kMulMatVectorPackedMaxOutputSize; + default: + return kMulMatVectorDenseMaxOutputSize; + } +} + +static bool vector_mul_mat_route_supports(const CommonMulMatMatch & match) { + return match.matched() && match.token_count == 1 && match.input_size >= 256 && match.input_size <= 32768 && + match.input_size % 32 == 0 && match.output_size <= vector_max_output_size(match.weight_format); +} + +static bool vector_postops_route_supports(const MulMatPostOpsMatch & match) { + return match.matched() && match.token_count == 1 && match.input_size >= 256 && match.input_size <= 32768 && + match.input_size % 32 == 0 && match.output_size <= vector_max_output_size(match.weight_format); +} + +static 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; + } + + 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(const VectorPublishPlan & plan) { + return plan.publication.format == CommonActivationPublicationFormat::F16Row ? GGML_TYPE_F16 : 0; +} + +static bool append_vector_q8_publish(const Value & output, + int64_t token_count, + int64_t output_size, + VectorPublishPlan & publish, + DispatchMatch & dispatch_match) { + if (publish.publication.format != CommonActivationPublicationFormat::Q8_1X4) { + return true; + } + + Dispatch quantize; + quantize.kernel = make_kernel_specialization(kQuantizeQ8_1X4F32Kernel); + quantize.kernel.integer_parameters.emplace("token_count", token_count); + quantize.kernel.integer_parameters.emplace("input_size", output_size); + quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", + common_to_config_value(token_count * output_size / 128)); + quantize.bindings.push_back({ output.id, 0, output.byte_count }); + quantize.bindings.push_back(publish.publication.binding()); + dispatch_match.dispatches.push_back(std::move(quantize)); + return true; +} + +static DispatchBinding vector_next_output_binding(const Value & output, const VectorPublishPlan & publish) { + if (publish.publication.format == CommonActivationPublicationFormat::F16Row) { + return publish.publication.binding(); + } + return { output.id, 0, output.byte_count }; +} + +static bool build_vector_mul_mat_dispatch(const DispatchMatchContext & context, + const CommonMulMatMatch & match, + DispatchMatch & dispatch_match) { + if (match.token_count != 1 || match.input_size % 32 != 0 || + match.output_size > vector_max_output_size(match.weight_format)) { + return false; + } + + if (match.input_size % 256 != 0 && is_k256_packed_row_layout(match.weight_format)) { + return false; + } + + Dispatch dispatch; + 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; + } + + 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; + const char * weight_layout = nullptr; + bool uses_q8_activation = false; + if (!prepare_vector_activation_binding(context, *match.input, *match.weight, match.input_size, match.output_size, + match.token_count, match.weight_format, dispatch_match, activation, + weight_format, weight_layout, uses_q8_activation)) { + return false; + } + + 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) { + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.activation_format", + std::to_string(GGML_TYPE_Q8_1)); + } + + dispatch.bindings.push_back(activation); + if (weight_layout != nullptr) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + if (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)); + + dispatch_match.covered_nodes.push_back(context.root_index); + if (match.has_fused_unary) { + dispatch_match.covered_nodes.push_back(match.unary_node_index); + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return append_vector_q8_publish(*match.output, match.token_count, match.output_size, publish, dispatch_match); +} + +static bool build_vector_mul_mat_postops_dispatch(const DispatchMatchContext & context, + const MulMatPostOpsMatch & match, + DispatchMatch & dispatch_match) { + if (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; + } + + const bool has_tail = match.input_size % 256 != 0; + if (has_tail && is_k256_packed_row_layout(match.weight_format)) { + return false; + } + + if (!match.has_bias && !match.has_residual) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatVectorBiasAddF32F32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.input_size", + common_to_config_value(match.input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.output_size", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector_postops.apply_bias", + common_to_config_value(match.has_bias ? int64_t{ 1 } : int64_t{ 0 })); + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector_postops.apply_residual", + common_to_config_value(match.has_residual ? int64_t{ 1 } : int64_t{ 0 })); + VectorPublishPlan publish; + 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))); + + DispatchBinding activation; + CommonMulMatWeightFormat weight_format = match.weight_format; + const char * weight_layout = nullptr; + bool uses_q8_activation = false; + if (!prepare_vector_activation_binding(context, *match.input, *match.weight, match.input_size, match.output_size, + match.token_count, match.weight_format, dispatch_match, activation, + weight_format, weight_layout, uses_q8_activation)) { + return false; + } + + dispatch.kernel.compile_parameters.emplace( + "ggml.matmul.vector_postops.weight_format", + common_to_config_value(common_mul_mat_format_config_value(weight_format))); + if (uses_q8_activation) { + dispatch.kernel.compile_parameters.emplace("ggml.matmul.vector.activation_format", + std::to_string(GGML_TYPE_Q8_1)); + } + + dispatch.bindings.push_back(activation); + if (weight_layout != nullptr) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, weight_layout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + const Value * placeholder = match.has_bias ? match.bias : match.residual_input; + dispatch.bindings.push_back(match.has_bias ? DispatchBinding{ match.bias->id, 0, match.bias->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back(match.has_residual ? + DispatchBinding{ match.residual_input->id, 0, match.residual_input->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + dispatch.bindings.push_back(vector_next_output_binding(*match.residual_output, publish)); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes)) { + return false; + } + if (match.layout_node != nullptr && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.layout_node, + dispatch_match.covered_nodes)) { + return false; + } + for (const GraphNode * add_node : match.add_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, add_node, + dispatch_match.covered_nodes)) { + return false; + } + } + + dispatch_match.dispatches.push_back(std::move(dispatch)); + return append_vector_q8_publish(*match.residual_output, match.token_count, match.output_size, publish, + dispatch_match); +} + +static bool match_vector_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatVectorF32F32Kernel, true); + if (!match.matched() || match.token_count != 1) { + return false; + } + return build_vector_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_vector_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatVectorF32F32Kernel, true); + if (!match.matched() || match.token_count != 1 || !try_match_fused_unary(context, match)) { + return false; + } + return build_vector_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_vector_mul_mat_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = match_mul_mat_postops(context, true); + if (!match.matched() || match.token_count != 1) { + return false; + } + return build_vector_mul_mat_postops_dispatch(context, match, dispatch_match); +} + +static bool match_q6_vector_final_projection_q8_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, true); + if (!match.matched()) { + match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32WmmaKernel, false); + } + if (!match.matched() || !context.graph.has_index() || match.weight->type != GGML_TYPE_Q6_K || + match.token_count != 1 || match.weight->alias_source.value >= 0 || match.input_size % 256 != 0 || + match.output_size % 64 != 0 || match.output_size > kMulMatVectorPackedMaxOutputSize || + !context.graph.index().consumers(match.output->id).empty()) { + return false; + } + + if (is_qwen_endpoint_rmsnorm_projection(context.graph, match)) { + return false; + } + + if (match.output_size < 16 * match.input_size) { + return false; + } + + return build_q6_vector_dispatch(match, context.root_index, dispatch_match); +} + +static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, + const MulMatPostOpsMatch & match, + DispatchMatch & dispatch_match) { + 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( + 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); + 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", + common_to_config_value(match.output_size)); + const bool pack_rows = match.token_count <= 5 && common_mul_mat_has_skinny_reduction_shape(match.input_size) && + match.output_size % 64 == 0 && match.weight->alias_source.value < 0; + const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; + const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; + bool pack_q4 = false; + bool pack_q6 = false; + DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; + if (match.requires_q8_activation && !can_pack_q4 && !can_pack_q6) { + return false; + } + if (can_pack_q4 || can_pack_q6) { + if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, + activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + pack_q4 = can_pack_q4; + pack_q6 = can_pack_q6; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", + std::to_string(GGML_TYPE_Q8_1)); + } + } + if (match.requires_q8_activation && !pack_q4 && !pack_q6) { + return false; + } + const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : + pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : + match.weight_format; + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_postops.weight_format", + common_to_config_value(common_mul_mat_format_config_value(format))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_postops.epilogue", common_to_config_value((match.has_bias ? int64_t{ 1 } : int64_t{ 0 }) | + (match.has_residual ? int64_t{ 2 } : int64_t{ 0 }))); + dispatch.bindings.push_back(activation); + if (pack_q4 || pack_q6) { + const char * layout = pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, layout, match.weight->type, + match.input_size, match.output_size, match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } + const Value * placeholder = match.has_bias ? match.bias : match.residual_input; + dispatch.bindings.push_back(match.has_bias ? DispatchBinding{ match.bias->id, 0, match.bias->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back(match.has_residual ? + DispatchBinding{ match.residual_input->id, 0, match.residual_input->byte_count } : + DispatchBinding{ placeholder->id, 0, placeholder->byte_count }); + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes)) { + return false; + } + if (match.layout_node != nullptr && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.layout_node, + dispatch_match.covered_nodes)) { + return false; + } + for (const GraphNode * add_node : match.add_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, add_node, + dispatch_match.covered_nodes)) { + return false; + } + } + + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_skinny_mul_mat_postops_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = + match_mul_mat_postops(context, kMulMatSkinnyF32F32Kernel, kMulMatSkinnyBiasAddF32F32Kernel, + kMulMatSkinnyBiasAddF32F32Kernel, kMulMatSkinnyBiasAddF32F32Kernel); + if (!common_mul_mat_postops_is_skinny_route(match)) { + return false; + } + return build_mul_mat_postops_dispatch(context, match, dispatch_match); +} + +static bool match_tiled_mul_mat_postops_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const MulMatPostOpsMatch match = + match_mul_mat_postops(context, kMulMatTiledF32F32Kernel, kMulMatTiledBiasAddF32F32Kernel, + kMulMatTiledBiasAddF32F32Kernel, kMulMatTiledBiasAddF32F32Kernel); + if (!common_mul_mat_postops_is_tiled_route(match)) { + return false; + } + return build_mul_mat_postops_dispatch(context, match, dispatch_match); +} + +static DecodeMulMatAddMatch match_decode_mul_mat_add(const DispatchMatchContext & context) { + DecodeMulMatAddMatch match; + CommonMulMatMatch root = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatAddF32F32DecodeWave64Kernel, true); + if (!root.matched() || !context.graph.has_index() || root.token_count != 1) { + return match; + } + + const GraphNode * add_node = common_find_only_consumer_with_op(context.graph, root.output->id, GGML_OP_ADD); + if (add_node == nullptr || !common_binary_node_is_add(*add_node)) { + return match; + } + + const bool root_is_lhs = add_node->inputs[0] == root.output->id; + const bool root_is_rhs = add_node->inputs[1] == root.output->id; + if (!root_is_lhs && !root_is_rhs) { + return match; + } + + const ValueId residual_id = root_is_lhs ? add_node->inputs[1] : add_node->inputs[0]; + const Value * residual = common_graph_value(context.graph, residual_id); + const Value * output = common_graph_value(context.graph, add_node->output); + if (residual == nullptr || output == nullptr || residual->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !residual->contiguous || !output->contiguous || !common_same_shape(*root.output, *residual) || + !common_same_shape(*root.output, *output)) { + return match; + } + + size_t add_index = 0; + if (!context.graph.index().node_index(add_node, add_index) || add_index >= context.covered_nodes.size() || + context.covered_nodes[add_index]) { + return match; + } + + root.output = output; + match.root = root; + match.residual_input = residual; + match.residual_output = output; + match.add_node = add_node; + return match; +} + +static bool build_decode_mul_mat_add_dispatch(const DecodeMulMatAddMatch & match, + const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kMulMatAddF32F32DecodeWave64Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.root.token_count); + dispatch.kernel.integer_parameters.emplace("input_size", match.root.input_size); + dispatch.kernel.integer_parameters.emplace("output_size", match.root.output_size); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_f32_f32_decode.token_capacity", + common_to_config_value(match.root.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_f32_f32_decode.output_capacity", + common_to_config_value(match.root.output_size)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_f32_f32_decode.weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.root.weight_format))); + dispatch.bindings.push_back({ match.root.input->id, 0, match.root.input->byte_count }); + dispatch.bindings.push_back({ match.root.weight->id, 0, match.root.weight->byte_count }); + dispatch.bindings.push_back({ match.residual_input->id, 0, match.residual_input->byte_count }); + dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.add_node, + dispatch_match.covered_nodes)) { + return false; + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static void build_decode_mul_mat_dispatch(const CommonMulMatMatch & match, + DispatchMatch & dispatch_match, + size_t root_index) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.integer_parameters.emplace("input_size", match.input_size); + dispatch.kernel.integer_parameters.emplace("output_size", match.output_size); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_f32_f32_decode.token_capacity", + common_to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_f32_f32_decode.output_capacity", + common_to_config_value(match.output_size)); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_f32_f32_decode.weight_format", + common_to_config_value(common_mul_mat_format_config_value(match.weight_format))); + dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); + dispatch.bindings.push_back({ 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(root_index); + dispatch_match.dispatches.push_back(std::move(dispatch)); +} + +static bool match_skinny_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSkinnyF32F32Kernel, false); + if (!match.matched()) { + return false; + } + try_match_fused_unary(context, match); + if (!common_mul_mat_is_skinny_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); +} + +static bool match_tiled_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledF32F32Kernel, false); + if (!match.matched()) { + return false; + } + try_match_fused_unary(context, match); + if (!common_mul_mat_is_tiled_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_skinny_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatSkinnyF32F32Kernel, false); + if (!match.matched() || !try_match_fused_unary(context, match) || !common_mul_mat_is_skinny_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match, CommonQ8ActivationPolicy::AllowStandaloneQuantize); +} + +static bool match_tiled_mul_mat_unary_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatTiledF32F32Kernel, false); + if (!match.matched() || !try_match_fused_unary(context, match) || !common_mul_mat_is_tiled_route(match)) { + return false; + } + return build_mul_mat_dispatch(context, match, dispatch_match); +} + +static bool match_decode_mul_mat_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + if (deprecated_mul_mat_dispatch_disabled()) { + return false; + } + CommonMulMatMatch match = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatF32F32DecodeWave64Kernel, true); + if (!match.matched()) { + return false; + } + if (iq_codebook_matmul_required(match.weight_format)) { + return false; + } + if (vector_mul_mat_route_supports(match)) { + return false; + } + if ((match.weight_format == CommonMulMatWeightFormat::Q4K || + match.weight_format == CommonMulMatWeightFormat::Q6K) && + match.output_size % 64 == 0 && common_is_supported_dense_output_size(match.output_size) && + match.weight->alias_source.value < 0) { + match.kernel = kMulMatF32F32WmmaKernel; + return build_mul_mat_dispatch(context, match, dispatch_match, + CommonQ8ActivationPolicy::AllowStandaloneQuantize); + } else { + build_decode_mul_mat_dispatch(match, dispatch_match, context.root_index); + } + return true; +} + +static bool match_decode_mul_mat_add_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + if (deprecated_mul_mat_dispatch_disabled()) { + return false; + } + const DecodeMulMatAddMatch match = match_decode_mul_mat_add(context); + if (!match.matched()) { + return false; + } + if (iq_codebook_matmul_required(match.root.weight_format)) { + return false; + } + MulMatPostOpsMatch postops; + postops.input = match.root.input; + postops.weight = match.root.weight; + postops.projection_output = match.root.output; + postops.residual_input = match.residual_input; + postops.residual_output = match.residual_output; + postops.kernel = kMulMatVectorBiasAddF32F32Kernel; + postops.weight_format = match.root.weight_format; + postops.input_size = match.root.input_size; + postops.output_size = match.root.output_size; + postops.token_count = match.root.token_count; + postops.has_residual = true; + if (vector_postops_route_supports(postops)) { + return false; + } + return build_decode_mul_mat_add_dispatch(match, context, dispatch_match); +} + +void register_mul_mat_dispatches(DispatchRegistryBuilder & registry) { + registry.add_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, + DispatchMatchKind::Fused, + 310, + DispatchSource::Common, + match_q6_vector_selected_outputs_dispatch, + }); + registry.add({ + "common.mul_mat.q6_k_prefill_wave32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 295, + DispatchSource::Common, + match_q6_k_prefill_wave32_dispatch, + }); + registry.add({ + "common.mul_mat.vector_q6_final_projection_q8", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 305, + DispatchSource::Common, + match_q6_vector_final_projection_q8_dispatch, + }); + registry.add({ + "common.mul_mat.vector_q6", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_q6_vector_dispatch, + }); + registry.add({ + "common.mul_mat.q5_k_iq4_xs_symmetric_i4_lowrow", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 285, + DispatchSource::Common, + match_symmetric_i4_low_row_dispatch, + }); + registry.add({ + "common.mul_mat.q5_k_iq4_xs_q8_1_x4_prefill", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 260, + DispatchSource::Common, + match_packed_q8_1_x4_prefill_dispatch, + }); + registry.add({ + "common.mul_mat.q4_k_q8_1_x4_prefill", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 305, + DispatchSource::Common, + match_q4_k_q8_1_x4_prefill_dispatch, + }); + registry.add({ + "common.mul_mat.q5_k_symmetric_i4_prefill", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 270, + DispatchSource::Common, + match_symmetric_i4_prefill_dispatch, + }); + registry.add({ + "common.mul_mat.q5_k_symmetric_i8_prefill", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 280, + DispatchSource::Common, + match_q5_k_symmetric_i8_prefill_dispatch, + }); + registry.add({ + "common.mul_mat_postops.vector", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 292, + DispatchSource::Common, + match_vector_mul_mat_postops_dispatch, + }); + registry.add({ + "common.mul_mat_unary.vector", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 291, + DispatchSource::Common, + match_vector_mul_mat_unary_dispatch, + }); + registry.add({ + "common.mul_mat_unary.tiled_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 291, + DispatchSource::Common, + match_tiled_mul_mat_unary_dispatch, + }); + registry.add({ + "common.mul_mat_unary.skinny_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 292, + DispatchSource::Common, + match_skinny_mul_mat_unary_dispatch, + }); + registry.add({ + "common.mul_mat.vector", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 190, + DispatchSource::Common, + match_vector_mul_mat_dispatch, + }); + registry.add({ + "common.mul_mat_postops.tiled_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 181, + DispatchSource::Common, + match_tiled_mul_mat_postops_dispatch, + }); + registry.add({ + "common.mul_mat_postops.skinny_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 182, + DispatchSource::Common, + match_skinny_mul_mat_postops_dispatch, + }); + registry.add({ + "common.mul_mat.tiled_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 81, + DispatchSource::Common, + match_tiled_mul_mat_dispatch, + }); + registry.add({ + "common.mul_mat.skinny_f32_f32", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 82, + DispatchSource::Common, + match_skinny_mul_mat_dispatch, + }); + registry.add({ + "common.mul_mat_add.skinny_f32_f32_decode", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 70, + DispatchSource::Common, + match_decode_mul_mat_add_dispatch, + }); + registry.add({ + "common.mul_mat.skinny_f32_f32_decode", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 60, + DispatchSource::Common, + match_decode_mul_mat_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.h new file mode 100644 index 000000000000..33a822434eaa --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_mul_mat_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp new file mode 100644 index 000000000000..741bd7c110f1 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.cpp @@ -0,0 +1,1215 @@ +#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" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kRmsNormBinaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32"); +static constexpr KernelCatalogRef kRmsNormBinaryF32LowTokenStripedKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_f32_lowtoken_striped"); +static constexpr KernelCatalogRef kRmsNormBinaryStridedF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_binary_strided_f32"); +static constexpr KernelCatalogRef kRmsNormMulAddF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rmsnorm_mul_add_f32"); +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_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"); +static constexpr KernelCatalogRef kAddRmsNormBinarySymmetricI4K32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_add_rmsnorm_binary_symmetric_i4_k32"); +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 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); +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_supported_f32_hidden_size(int64_t hidden_size) { + return hidden_size >= 64 && hidden_size <= 32768 && hidden_size % 64 == 0; +} + +static bool is_supported_binary_f32_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +static bool is_supported_symmetric_i4_hidden_size(int64_t hidden_size) { + return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; +} + +static bool is_supported_token_count(int64_t token_count) { + return token_count >= 1 && token_count <= (1 << 20); +} + +static bool supported_rmsnorm_input_layout(const Value & input, + int64_t hidden_size, + int64_t token_count, + int64_t & input_stride, + size_t & input_span_bytes) { + if (input.type != GGML_TYPE_F32 || input.ne[0] != hidden_size || input.nb[0] != sizeof(float)) { + return false; + } + + if (input.contiguous) { + input_stride = hidden_size; + input_span_bytes = input.byte_count; + return true; + } + + if (input.ne[1] != token_count || input.ne[2] != 1 || input.ne[3] != 1 || input.nb[1] % sizeof(float) != 0) { + return false; + } + input_stride = static_cast(input.nb[1] / sizeof(float)); + if (input_stride < hidden_size || input_stride > 1048576) { + return false; + } + + return strided_f32_storage_span_bytes(input, input_span_bytes); +} + +static bool is_weight_shape(const Value & weight, int64_t hidden_size) { + if (weight.ne[0] != hidden_size) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (weight.ne[i] != 1) { + return false; + } + } + return true; +} + +struct RmsNormBinaryMatch { + const GraphNode * rms_node = nullptr; + const GraphNode * binary_node = nullptr; + const Value * input = nullptr; + const Value * rhs = nullptr; + const Value * output = nullptr; + size_t rms_node_index = 0; + size_t binary_node_index = 0; + int64_t hidden_size = 0; + int64_t token_count = 0; + int64_t input_stride = 0; + size_t input_span = 0; + BinaryKind op = BinaryKind::Add; + float epsilon = 0.0f; + + bool matched() const { + return rms_node != nullptr && binary_node != nullptr && input != nullptr && rhs != nullptr && output != nullptr; + } +}; + +struct RmsNormMulAddMatch { + RmsNormBinaryMatch rms_binary; + const GraphNode * add_node = nullptr; + const Value * residual = nullptr; + const Value * output = nullptr; + size_t add_node_index = 0; + BinaryKind output_op = BinaryKind::Add; + bool binary_lhs = true; + + bool matched() const { + return rms_binary.matched() && add_node != nullptr && residual != nullptr && output != nullptr; + } +}; + +struct RmsNormMatch { + const GraphNode * rms_node = nullptr; + const Value * input = nullptr; + const Value * output = nullptr; + size_t input_span = 0; + size_t rms_node_index = 0; + int64_t input_stride = 0; + int64_t hidden_size = 0; + int64_t token_count = 0; + float epsilon = 0.0f; + + bool matched() const { return rms_node != nullptr && input != nullptr && output != nullptr; } +}; + +struct RmsNormMulRopeMatch { + const GraphNode * rms = nullptr; + const GraphNode * binary = nullptr; + const GraphNode * rope = nullptr; + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * positions = nullptr; + const Value * output = nullptr; + const RmsNormParams * rms_params = nullptr; + const RopeParams * rope_params = nullptr; + + bool matched() const { + return rms != nullptr && binary != nullptr && rope != nullptr && input != nullptr && weight != nullptr && + positions != nullptr && output != nullptr && rms_params != nullptr && rope_params != nullptr; + } +}; + +struct AddRmsNormBinarySymmetricI4Match { + const GraphNode * add_node = nullptr; + const Value * lhs = nullptr; + const Value * rhs = nullptr; + const Value * residual = nullptr; + RmsNormBinaryMatch rms_binary; + size_t add_node_index = 0; + + bool matched() const { + return add_node != nullptr && lhs != nullptr && rhs != nullptr && residual != nullptr && rms_binary.matched(); + } +}; + +struct RmsNormGateMatch { + std::vector covered; + const Value * input = nullptr; + const Value * weight = nullptr; + const Value * raw_gate = nullptr; + const Value * output = nullptr; + int64_t hidden_size = 0; + int64_t token_count = 0; + float epsilon = 0.0f; + UnaryKind gate_op = UnaryKind::Silu; + + bool matched() const { + return !covered.empty() && input != nullptr && weight != nullptr && raw_gate != nullptr && output != nullptr; + } +}; + +template static bool pairwise_distinct_storage_roots(const std::array & values) { + for (size_t lhs = 0; lhs < values.size(); ++lhs) { + for (size_t rhs = lhs + 1; rhs < values.size(); ++rhs) { + if (values[lhs]->storage_root == values[rhs]->storage_root) { + return false; + } + } + } + return true; +} + +static bool rope_has_set_rows_consumer(const Graph & graph, const GraphNode & rope) { + std::vector pending = { rope.output }; + for (size_t i = 0; i < pending.size(); ++i) { + for (const GraphNode * consumer : graph.index().consumers(pending[i])) { + if (consumer == nullptr) { + continue; + } + if (consumer->op == GGML_OP_SET_ROWS) { + return true; + } + if (is_layout_alias_node(graph, *consumer)) { + pending.push_back(consumer->output); + } + } + } + return false; +} + +static bool fused_input_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +static RmsNormMulRopeMatch match_rmsnorm_mul_rope_f32(const Graph & graph, const GraphNode * rms) { + RmsNormMulRopeMatch match; + if (rms == nullptr || rms->op != GGML_OP_RMS_NORM || rms->inputs.size() != 1 || !graph.has_index()) { + return match; + } + + const std::vector & rms_consumers = graph.index().consumers(rms->output); + if (rms_consumers.size() != 1 || rms_consumers.front() == nullptr || rms_consumers.front()->op != GGML_OP_MUL || + rms_consumers.front()->inputs.size() != 2) { + return {}; + } + const GraphNode * binary = rms_consumers.front(); + const std::vector & binary_consumers = graph.index().consumers(binary->output); + if (binary_consumers.size() != 1 || binary_consumers.front() == nullptr || + binary_consumers.front()->op != GGML_OP_ROPE || binary_consumers.front()->inputs.size() != 2 || + binary_consumers.front()->inputs[0] != binary->output) { + return {}; + } + const GraphNode * rope = binary_consumers.front(); + + const Value * input = graph_value(graph, rms->inputs[0]); + const Value * rms_output = graph_value(graph, rms->output); + const Value * binary_output = graph_value(graph, binary->output); + const Value * positions = graph_value(graph, rope->inputs[1]); + const Value * output = graph_value(graph, rope->output); + const Value * weight = nullptr; + if (binary->inputs[0] == rms->output) { + weight = graph_value(graph, binary->inputs[1]); + } else if (binary->inputs[1] == rms->output) { + weight = graph_value(graph, binary->inputs[0]); + } + const RmsNormParams * rms_params = op_params_as(rms->params); + const RopeParams * rope_params = op_params_as(rope->params); + if (input == nullptr || rms_output == nullptr || binary_output == nullptr || positions == nullptr || + output == nullptr || weight == nullptr || rms_params == nullptr || rope_params == nullptr) { + return {}; + } + + if (input->type != GGML_TYPE_F32 || rms_output->type != GGML_TYPE_F32 || binary_output->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || weight->type != GGML_TYPE_F32 || positions->type != GGML_TYPE_I32 || + input->ne != rms_output->ne || input->ne != binary_output->ne || input->ne != output->ne || + !rms_output->contiguous || !binary_output->contiguous || !output->contiguous || !weight->contiguous || + !positions->contiguous || weight->ne[0] != input->ne[0] || weight->ne[1] != 1 || weight->ne[2] != 1 || + weight->ne[3] != 1 || input->ne[0] < 2 || input->ne[0] > 512 || input->ne[0] % 2 != 0 || input->ne[1] < 1 || + input->ne[1] > 1048576 || input->ne[2] < 1 || input->ne[2] > 1048576 || input->ne[3] < 1 || + input->ne[3] > 1024 || input->nb[0] != sizeof(float) || output->nb[0] != sizeof(float)) { + return {}; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (input->nb[i] % sizeof(float) != 0 || output->nb[i] % sizeof(float) != 0) { + return {}; + } + } + + const bool sectioned_mode = rope_params->mode == GGML_ROPE_TYPE_MROPE || rope_params->mode == GGML_ROPE_TYPE_IMROPE; + const bool lfm_neox_mode = rope_params->mode == GGML_ROPE_TYPE_NEOX && input->ne[0] == 64 && + rope_params->n_dims == 64 && !rope_has_set_rows_consumer(graph, *rope); + int64_t section_count = 0; + if (sectioned_mode) { + for (int section : rope_params->sections) { + if (section < 0) { + return {}; + } + section_count += section; + } + } + const int64_t required_positions = sectioned_mode ? 4 * input->ne[2] : input->ne[2]; + if ((!sectioned_mode && !lfm_neox_mode) || positions->element_count < required_positions || + rope_params->n_dims < 2 || rope_params->n_dims > input->ne[0] || rope_params->n_dims % 2 != 0 || + (sectioned_mode && section_count != rope_params->n_dims / 2) || !std::isfinite(rms_params->eps) || + rms_params->eps <= 0.0f || !std::isfinite(rope_params->freq_base) || rope_params->freq_base <= 0.0f || + !std::isfinite(rope_params->freq_scale) || rope_params->freq_scale <= 0.0f || + !std::isfinite(rope_params->attn_factor) || rope_params->attn_factor <= 0.0f || + rope_params->ext_factor != 0.0f || output->storage_root == input->storage_root || + output->storage_root == weight->storage_root || output->storage_root == positions->storage_root) { + return {}; + } + + match.rms = rms; + match.binary = binary; + match.rope = rope; + match.input = input; + match.weight = weight; + match.positions = positions; + match.output = output; + match.rms_params = rms_params; + match.rope_params = rope_params; + return match; +} + +static RmsNormMatch match_rmsnorm_f32(const Graph & graph, const GraphNode * node, size_t node_index) { + RmsNormMatch match; + if (node == nullptr || node->op != GGML_OP_RMS_NORM || node->inputs.size() != 1) { + return match; + } + + const RmsNormParams * rms_params = op_params_as(node->params); + if (rms_params == nullptr || !std::isfinite(rms_params->eps) || rms_params->eps <= 0.0f) { + return {}; + } + + const Value * input = graph_value(graph, node->inputs[0]); + const Value * output = graph_value(graph, node->output); + if (input == nullptr || output == nullptr) { + return {}; + } + if (input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || !output->contiguous || + !same_shape(*input, *output)) { + return {}; + } + if (input->storage_root == output->storage_root) { + return {}; + } + + const int64_t hidden_size = output->ne[0]; + if (!is_supported_f32_hidden_size(hidden_size)) { + return {}; + } + if (hidden_size == 0 || output->element_count <= 0 || output->element_count % hidden_size != 0) { + return {}; + } + const int64_t token_count = output->element_count / hidden_size; + if (!is_supported_token_count(token_count)) { + return {}; + } + int64_t input_stride = 0; + size_t input_span = 0; + if (!supported_rmsnorm_input_layout(*input, hidden_size, token_count, input_stride, input_span)) { + return {}; + } + + match.rms_node = node; + match.input = input; + match.output = output; + match.input_span = input_span; + match.rms_node_index = node_index; + match.input_stride = input_stride; + match.hidden_size = hidden_size; + match.token_count = token_count; + match.epsilon = rms_params->eps; + return match; +} + +static RmsNormBinaryMatch match_rmsnorm_binary_f32(const Graph & graph, + const GraphNode * node, + size_t node_index, + bool allow_strided_input = false) { + RmsNormBinaryMatch match; + if (node == nullptr || node->op != GGML_OP_RMS_NORM || node->inputs.size() != 1 || !graph.has_index()) { + return match; + } + + const RmsNormParams * rms_params = op_params_as(node->params); + if (rms_params == nullptr || !std::isfinite(rms_params->eps) || rms_params->eps <= 0.0f) { + return {}; + } + const std::vector & consumers = graph.index().consumers(node->output); + if (consumers.size() != 1) { + return {}; + } + const GraphNode * binary_node = consumers.front(); + size_t binary_node_index; + const CommonBinaryMatch binary_match = + common_match_binary_consumer(binary_node, node->output, kCommonArithmeticBinaryKinds); + if (!binary_match.matched() || !graph.index().node_index(binary_node, binary_node_index)) { + return {}; + } + if (!binary_match.producer_is_lhs && common_binary_kind_requires_order(binary_match.kind)) { + return {}; + } + + const Value * input = graph_value(graph, node->inputs[0]); + const Value * rms = graph_value(graph, node->output); + const Value * rhs = graph_value(graph, binary_match.operand); + const Value * output = graph_value(graph, binary_node->output); + if (input == nullptr || rms == nullptr || rhs == nullptr || output == nullptr) { + return {}; + } + if (input->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || rhs->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32) { + return {}; + } + if (!rms->contiguous || !rhs->contiguous || !output->contiguous) { + return {}; + } + if (!same_shape(*input, *rms) || !same_shape(*input, *output)) { + return {}; + } + + const int64_t hidden_size = output->ne[0]; + if (!is_supported_binary_f32_hidden_size(hidden_size) || !is_weight_shape(*rhs, hidden_size)) { + return {}; + } + if (hidden_size == 0 || output->element_count <= 0 || output->element_count % hidden_size != 0) { + return {}; + } + const int64_t token_count = output->element_count / hidden_size; + if (!is_supported_token_count(token_count)) { + return {}; + } + int64_t input_stride = 0; + size_t input_span = 0; + if (!supported_rmsnorm_input_layout(*input, hidden_size, token_count, input_stride, input_span) || + (!allow_strided_input && input_stride != hidden_size) || + !pairwise_distinct_storage_roots(std::array{ input, rhs, output })) { + return {}; + } + + match.rms_node = node; + match.binary_node = binary_node; + match.input = input; + match.rhs = rhs; + match.output = output; + match.rms_node_index = node_index; + match.binary_node_index = binary_node_index; + match.hidden_size = hidden_size; + match.token_count = token_count; + match.input_stride = input_stride; + match.input_span = input_span; + match.op = binary_match.kind; + match.epsilon = rms_params->eps; + return match; +} + +static RmsNormMulAddMatch match_rmsnorm_mul_add_f32(const DispatchMatchContext & context) { + RmsNormMulAddMatch match; + match.rms_binary = match_rmsnorm_binary_f32(context.graph, context.root_node, context.root_index); + if (!match.rms_binary.matched()) { + return {}; + } + + const std::vector & consumers = context.graph.index().consumers(match.rms_binary.output->id); + if (consumers.size() != 1) { + return {}; + } + const GraphNode * add = consumers.front(); + size_t add_index = 0; + const CommonBinaryMatch output_binary = + common_match_binary_consumer(add, match.rms_binary.output->id, kCommonArithmeticBinaryKinds); + if (!output_binary.matched() || !context.graph.index().node_index(add, add_index)) { + return {}; + } + const Value * residual = graph_value(context.graph, output_binary.operand); + const Value * output = graph_value(context.graph, add->output); + const Value * rms_output = graph_value(context.graph, match.rms_binary.rms_node->output); + if (residual == nullptr || output == nullptr || rms_output == nullptr || residual->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !residual->contiguous || !output->contiguous || + !same_shape(*residual, *match.rms_binary.input) || !same_shape(*output, *match.rms_binary.input) || + !pairwise_distinct_storage_roots(std::array{ + match.rms_binary.input, match.rms_binary.rhs, rms_output, match.rms_binary.output, residual, output })) { + return {}; + } + const size_t nodes_size = context.covered_nodes.size(); + if (match.rms_binary.rms_node_index >= nodes_size || match.rms_binary.binary_node_index >= nodes_size || + add_index >= nodes_size || context.covered_nodes[match.rms_binary.rms_node_index] || + context.covered_nodes[match.rms_binary.binary_node_index] || context.covered_nodes[add_index]) { + return {}; + } + const auto available = [&](const Value & value) { + const GraphNode * producer = context.graph.index().producer(value.id); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return context.graph.index().node_index(producer, producer_index) && producer_index < nodes_size && + context.covered_nodes[producer_index]; + }; + if (!available(*match.rms_binary.input) || !available(*match.rms_binary.rhs) || !available(*residual)) { + return {}; + } + + match.add_node = add; + match.residual = residual; + match.output = output; + match.add_node_index = add_index; + match.output_op = output_binary.kind; + match.binary_lhs = output_binary.producer_is_lhs; + return match; +} + +static AddRmsNormBinarySymmetricI4Match match_add_rmsnorm_binary_symmetric_i4(const Graph & graph, + const GraphNode * node, + size_t node_index) { + AddRmsNormBinarySymmetricI4Match match; + if (node == nullptr || node->op != GGML_OP_ADD || node->inputs.size() != 2 || !graph.has_index()) { + return match; + } + + const Value * lhs = graph_value(graph, node->inputs[0]); + const Value * rhs = graph_value(graph, node->inputs[1]); + const Value * residual = graph_value(graph, node->output); + if (lhs == nullptr || rhs == nullptr || residual == nullptr || lhs->type != GGML_TYPE_F32 || + rhs->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || !lhs->contiguous || !rhs->contiguous || + !residual->contiguous || !same_shape(*lhs, *residual) || !same_shape(*rhs, *residual)) { + return {}; + } + + const GraphNode * rms_node = nullptr; + for (const GraphNode * consumer : graph.index().consumers(residual->id)) { + if (consumer == nullptr || consumer->op != GGML_OP_RMS_NORM) { + continue; + } + if (rms_node != nullptr) { + return {}; + } + rms_node = consumer; + } + size_t rms_node_index = 0; + if (rms_node == nullptr || !graph.index().node_index(rms_node, rms_node_index)) { + return {}; + } + + 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) || + !pairwise_distinct_storage_roots( + std::array{ lhs, rhs, residual, rms_binary.rhs, rms_binary.output })) { + return {}; + } + + match.add_node = node; + match.lhs = lhs; + match.rhs = rhs; + match.residual = residual; + match.rms_binary = rms_binary; + match.add_node_index = node_index; + return match; +} + +static RmsNormGateMatch match_rmsnorm_gate(const Graph & graph, const GraphNode * node, size_t node_index) { + RmsNormGateMatch match; + const RmsNormBinaryMatch rms_binary = match_rmsnorm_binary_f32(graph, node, node_index); + if (!rms_binary.matched() || rms_binary.op != BinaryKind::Mul || rms_binary.hidden_size % 64 != 0) { + return match; + } + + const GraphNode * terminal = common_find_only_consumer_with_op(graph, rms_binary.output->id, GGML_OP_MUL); + if (terminal == nullptr || !common_binary_node_is_mul(*terminal)) { + return {}; + } + const ValueId activated_id = + terminal->inputs[0] == rms_binary.output->id ? terminal->inputs[1] : terminal->inputs[0]; + const Value * activated = graph_value(graph, activated_id); + const GraphNode * unary = activated != nullptr ? graph.index().producer(activated->id) : nullptr; + const UnaryParams * unary_params = unary != nullptr ? op_params_as(unary->params) : nullptr; + if (unary == nullptr || unary->op != GGML_OP_UNARY || unary->inputs.size() != 1 || unary_params == nullptr || + !unary_kind_supported(unary_params->op) || + common_find_only_consumer_with_op(graph, activated->id, GGML_OP_MUL) != terminal) { + return {}; + } + + const Value * gate_input = graph_value(graph, unary->inputs[0]); + const GraphNode * gate_reshape = gate_input != nullptr ? graph.index().producer(gate_input->id) : nullptr; + const bool has_gate_reshape = + gate_reshape != nullptr && gate_reshape->op == GGML_OP_RESHAPE && gate_reshape->inputs.size() == 1; + const Value * raw_gate = has_gate_reshape ? graph_value(graph, gate_reshape->inputs[0]) : gate_input; + const Value * output = graph_value(graph, terminal->output); + if (gate_input == nullptr || raw_gate == nullptr || output == nullptr || gate_input->type != GGML_TYPE_F32 || + activated->type != GGML_TYPE_F32 || raw_gate->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !gate_input->contiguous || !activated->contiguous || !raw_gate->contiguous || !output->contiguous || + !same_shape(*rms_binary.input, *gate_input) || !same_shape(*rms_binary.input, *activated) || + !same_shape(*rms_binary.input, *output) || raw_gate->element_count != output->element_count || + graph.index().consumers(gate_input->id).size() != 1 || + !pairwise_distinct_storage_roots( + std::array{ rms_binary.input, rms_binary.rhs, raw_gate, output })) { + return {}; + } + + match.covered = { rms_binary.rms_node, rms_binary.binary_node }; + if (has_gate_reshape) { + match.covered.push_back(gate_reshape); + } + match.covered.push_back(unary); + match.covered.push_back(terminal); + match.input = rms_binary.input; + match.weight = rms_binary.rhs; + match.raw_gate = raw_gate; + match.output = output; + match.hidden_size = rms_binary.hidden_size; + match.token_count = rms_binary.token_count; + match.epsilon = rms_binary.epsilon; + match.gate_op = unary_params->op; + return match; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static std::string to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +} // namespace + +static bool match_rmsnorm_mul_rope_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const RmsNormMulRopeMatch fused = match_rmsnorm_mul_rope_f32(context.graph, context.root_node); + if (!fused.matched()) { + return false; + } + if (fused.rope_params->mode == GGML_ROPE_TYPE_NEOX && + (!fused_input_is_available(context.graph, fused.input->id, context.covered_nodes) || + !fused_input_is_available(context.graph, fused.weight->id, context.covered_nodes) || + !fused_input_is_available(context.graph, fused.positions->id, context.covered_nodes))) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRmsNormMulRopeF32Kernel); + auto & config = dispatch.kernel.compile_parameters; + config.emplace("ggml.rmsnorm_mul_rope.hidden_size", to_config_value(fused.input->ne[0])); + config.emplace("ggml.rmsnorm_mul_rope.ne1", to_config_value(fused.input->ne[1])); + config.emplace("ggml.rmsnorm_mul_rope.ne2", to_config_value(fused.input->ne[2])); + config.emplace("ggml.rmsnorm_mul_rope.ne3", to_config_value(fused.input->ne[3])); + config.emplace("ggml.rmsnorm_mul_rope.input_stride1", + to_config_value(static_cast(fused.input->nb[1] / sizeof(float)))); + config.emplace("ggml.rmsnorm_mul_rope.input_stride2", + to_config_value(static_cast(fused.input->nb[2] / sizeof(float)))); + config.emplace("ggml.rmsnorm_mul_rope.input_stride3", + to_config_value(static_cast(fused.input->nb[3] / sizeof(float)))); + config.emplace("ggml.rmsnorm_mul_rope.output_stride1", + to_config_value(static_cast(fused.output->nb[1] / sizeof(float)))); + config.emplace("ggml.rmsnorm_mul_rope.output_stride2", + to_config_value(static_cast(fused.output->nb[2] / sizeof(float)))); + config.emplace("ggml.rmsnorm_mul_rope.output_stride3", + to_config_value(static_cast(fused.output->nb[3] / sizeof(float)))); + config.emplace("ggml.rmsnorm_mul_rope.n_dims", to_config_value(static_cast(fused.rope_params->n_dims))); + const bool neox_mode = fused.rope_params->mode == GGML_ROPE_TYPE_NEOX; + config.emplace("ggml.rmsnorm_mul_rope.section0", + to_config_value(neox_mode ? static_cast(fused.rope_params->n_dims / 2) : + static_cast(fused.rope_params->sections[0]))); + config.emplace("ggml.rmsnorm_mul_rope.section1", + to_config_value(static_cast(fused.rope_params->sections[1]))); + config.emplace("ggml.rmsnorm_mul_rope.section2", + to_config_value(static_cast(fused.rope_params->sections[2]))); + config.emplace("ggml.rmsnorm_mul_rope.section3", + to_config_value(static_cast(fused.rope_params->sections[3]))); + config.emplace("ggml.rmsnorm_mul_rope.mode", to_config_value(static_cast(fused.rope_params->mode))); + config.emplace("ggml.rmsnorm_mul_rope.workgroup_size", "256"); + config.emplace("ggml.rmsnorm_mul_rope.epsilon", to_config_value(fused.rms_params->eps)); + config.emplace("ggml.rmsnorm_mul_rope.freq_base", to_config_value(fused.rope_params->freq_base)); + config.emplace("ggml.rmsnorm_mul_rope.freq_scale", to_config_value(fused.rope_params->freq_scale)); + config.emplace("ggml.rmsnorm_mul_rope.attn_factor", to_config_value(fused.rope_params->attn_factor)); + 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.positions->id, 0, fused.positions->byte_count }); + dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + if (!append_covered_node_index_once(context.graph, context.covered_nodes, fused.binary, match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, fused.rope, match.covered_nodes)) { + return false; + } + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_add_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchContext & context, + DispatchMatch & match) { + const AddRmsNormBinarySymmetricI4Match fused = + match_add_rmsnorm_binary_symmetric_i4(context.graph, context.root_node, context.root_index); + if (!fused.matched() || fused.add_node_index >= context.covered_nodes.size() || + fused.rms_binary.rms_node_index >= context.covered_nodes.size() || + fused.rms_binary.binary_node_index >= context.covered_nodes.size() || + context.covered_nodes[fused.add_node_index] || context.covered_nodes[fused.rms_binary.rms_node_index] || + context.covered_nodes[fused.rms_binary.binary_node_index]) { + return false; + } + + 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; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kAddRmsNormBinarySymmetricI4K32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", fused.rms_binary.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.add_rmsnorm_binary_symmetric_i4.hidden_size", + to_config_value(fused.rms_binary.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon", + to_config_value(fused.rms_binary.epsilon)); + dispatch.bindings.push_back({ fused.lhs->id, 0, fused.lhs->byte_count }); + dispatch.bindings.push_back({ fused.rhs->id, 0, fused.rhs->byte_count }); + 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(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)); + return match.status.success(); +} + +static bool match_rmsnorm_binary_symmetric_i4_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + const RmsNormBinaryMatch fused = + match_rmsnorm_binary_f32(context.graph, &nodes[context.root_index], context.root_index); + if (!fused.matched() || fused.op != BinaryKind::Mul || !is_supported_symmetric_i4_hidden_size(fused.hidden_size) || + 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] || + !pairwise_distinct_storage_roots(std::array{ fused.input, fused.rhs, fused.output })) { + return false; + } + + 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; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRmsNormBinarySymmetricI4K32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", fused.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_symmetric_i4.hidden_size", + to_config_value(fused.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_symmetric_i4.rms_epsilon", + to_config_value(fused.epsilon)); + 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(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)); + return match.status.success(); +} + +static bool match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const RmsNormGateMatch fused = match_rmsnorm_gate(context.graph, context.root_node, context.root_index); + if (!fused.matched()) { + return false; + } + + 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; + } + const Value * packed_input = fused.output; + const GraphNode * consumer = + common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_MUL_MAT); + if (consumer == nullptr) { + const GraphNode * reshape = + common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_RESHAPE); + const Value * reshaped = reshape != nullptr ? graph_value(context.graph, reshape->output) : nullptr; + if (reshaped != nullptr && is_layout_alias_node(context.graph, *reshape) && reshaped->contiguous && + same_full_value_range(*packed_input, *reshaped)) { + packed_input = reshaped; + consumer = common_find_only_consumer_with_op(context.graph, packed_input->id, GGML_OP_MUL_MAT); + } + } + 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"; + 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(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)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_f32.rms_epsilon", to_config_value(fused.epsilon)); + 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 ? 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(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.dispatches.push_back(std::move(dispatch)); + return true; + } + + 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; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRmsNormGateSiluMulSymmetricI4K32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", fused.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size", + to_config_value(fused.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon", + to_config_value(fused.epsilon)); + 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(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)); + return match.status.success(); +} + +static bool rmsnorm_k16_publication_enabled(const RmsNormBinaryMatch & rms) { + const bool qualified_shape = + (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); + 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 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) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + const RmsNormBinaryMatch rms_match = + match_rmsnorm_binary_f32(context.graph, &nodes[context.root_index], context.root_index, true); + 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]) { + return false; + } + + const bool strided_input = rms_match.input_stride != rms_match.hidden_size; + 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; + const bool lowtoken_striped = !strided_input && publication_demand == nullptr && rms_match.op == BinaryKind::Mul && + rms_match.token_count < 6 && rms_match.hidden_size >= 2048 && + rms_match.hidden_size <= 5120; + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(strided_input ? kRmsNormBinaryStridedF32Kernel : + packed_output ? kRmsNormBinaryF32K16Kernel : + f16_output ? kRmsNormBinaryF32F16Kernel : + lowtoken_striped ? kRmsNormBinaryF32LowTokenStripedKernel : + kRmsNormBinaryF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.hidden_size", + to_config_value(rms_match.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.rms_epsilon", + to_config_value(rms_match.epsilon)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.op", + std::to_string(binary_kind_config_value(rms_match.op))); + if (strided_input) { + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_f32.input_stride", + to_config_value(rms_match.input_stride)); + } + dispatch.bindings.push_back( + strided_input ? + DispatchBinding{ rms_match.input->storage_root, rms_match.input->storage_offset, rms_match.input_span } : + DispatchBinding{ 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 }); + + 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(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 match.status.success(); +} + +static bool match_rmsnorm_mul_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const RmsNormMulAddMatch fused = match_rmsnorm_mul_add_f32(context); + if (!fused.matched()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRmsNormMulAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", fused.rms_binary.token_count); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_mul_add_f32.hidden_size", + to_config_value(fused.rms_binary.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_mul_add_f32.rms_epsilon", + to_config_value(fused.rms_binary.epsilon)); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.rmsnorm_mul_add_f32.normalized_op", nullptr, + fused.rms_binary.op, true); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.rmsnorm_mul_add_f32.output_op", + "ggml.rmsnorm_mul_add_f32.binary_lhs", fused.output_op, fused.binary_lhs); + dispatch.bindings.push_back({ fused.rms_binary.input->id, 0, fused.rms_binary.input->byte_count }); + dispatch.bindings.push_back({ fused.rms_binary.rhs->id, 0, fused.rms_binary.rhs->byte_count }); + dispatch.bindings.push_back({ fused.residual->id, 0, fused.residual->byte_count }); + dispatch.bindings.push_back({ fused.output->id, 0, fused.output->byte_count }); + match.covered_nodes.push_back(fused.rms_binary.rms_node_index); + match.covered_nodes.push_back(fused.rms_binary.binary_node_index); + match.covered_nodes.push_back(fused.add_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_rmsnorm_binary_q8_1_x4_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + 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] || + q8_demand == nullptr || f16_demand == nullptr) { + return false; + } + + 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; + } + + Dispatch dispatch; + 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)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_binary_q8_1_x4.rms_epsilon", + 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_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)); + return match.status.success(); +} + +static bool match_rmsnorm_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const std::vector & nodes = context.graph.nodes(); + if (context.root_index >= nodes.size()) { + return false; + } + const RmsNormMatch rms_match = match_rmsnorm_f32(context.graph, &nodes[context.root_index], context.root_index); + if (!rms_match.matched() || rms_match.rms_node_index >= context.covered_nodes.size() || + context.covered_nodes[rms_match.rms_node_index]) { + return false; + } + + 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)); + dispatch.kernel.compile_parameters.emplace("ggml.rmsnorm_f32.input_stride", + to_config_value(rms_match.input_stride)); + dispatch.bindings.push_back( + { rms_match.input->storage_root, rms_match.input->storage_offset, rms_match.input_span }); + dispatch.bindings.push_back({ rms_match.output->id, 0, rms_match.output->byte_count }); + + match.covered_nodes.push_back(rms_match.rms_node_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +bool common_match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + return match_rmsnorm_gate_dispatch(context, match); +} + +void register_rmsnorm_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.add_rmsnorm_binary_symmetric_i4_k32", + GGML_OP_ADD, + DispatchMatchKind::Fused, + 400, + DispatchSource::Common, + match_add_rmsnorm_binary_symmetric_i4_dispatch, + }); + registry.add({ + "common.rmsnorm_gate", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 400, + DispatchSource::Common, + match_rmsnorm_gate_dispatch, + }); + registry.add({ + "common.rmsnorm_binary_symmetric_i4_k32", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 350, + DispatchSource::Common, + match_rmsnorm_binary_symmetric_i4_dispatch, + }); + registry.add({ + "common.rmsnorm_mul_rope_f32", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_rmsnorm_mul_rope_f32_dispatch, + }); + registry.add({ + "common.rmsnorm_mul_add_f32", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 250, + DispatchSource::Common, + match_rmsnorm_mul_add_f32_dispatch, + }); + registry.add({ + "common.rmsnorm_binary_q8_1_x4", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 150, + DispatchSource::Common, + match_rmsnorm_binary_q8_1_x4_dispatch, + }); + registry.add({ + "common.rmsnorm_binary_f32", + GGML_OP_RMS_NORM, + DispatchMatchKind::Fused, + 100, + DispatchSource::Common, + match_rmsnorm_binary_f32_dispatch, + }); + registry.add({ + "common.rmsnorm_f32", + GGML_OP_RMS_NORM, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_rmsnorm_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.h new file mode 100644 index 000000000000..12d043813715 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rmsnorm.h @@ -0,0 +1,11 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +bool common_match_rmsnorm_gate_dispatch(const DispatchMatchContext & context, DispatchMatch & match); + +void register_rmsnorm_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..3b668b05bb28 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.cpp @@ -0,0 +1,901 @@ +#include "dispatch-rope-set-rows.h" + +#include "dispatch-layout-utils.h" +#include "dispatch-rope-utils.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kRopeF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_f32"); +static constexpr KernelCatalogRef kRopeConcatF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_concat_f32"); +static constexpr KernelCatalogRef kRopeScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_scale_f32"); +static constexpr KernelCatalogRef kRopeTokenScaleF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_token_scale_f32"); +static constexpr KernelCatalogRef kSetRowsKernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_set_rows"); +static constexpr KernelCatalogRef kRopeSetRowsF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_rope_set_rows_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool is_1d_shape(const Value & value, int64_t ne0) { + return value.ne[0] == ne0 && value.ne[1] == 1 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_2d_shape(const Value & value, int64_t ne0, int64_t ne1) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static std::string value_layout_for_log(const Value * value) { + if (value == nullptr) { + return "null"; + } + std::ostringstream out; + out << ggml_type_name(value->type) << " ne=[" << value->ne[0] << "," << value->ne[1] << "," << value->ne[2] << "," + << value->ne[3] << "] nb=[" << value->nb[0] << "," << value->nb[1] << "," << value->nb[2] << "," << value->nb[3] + << "] contiguous=" << (value->contiguous ? 1 : 0); + if (value->alias_source.value >= 0) { + out << " alias=" << value->alias_source.value << " storage_offset=" << value->storage_offset; + } + return out.str(); +} + +static void log_rope_reject(Status * status, + const char * reason, + const Value * input, + const Value * positions, + const Value * output, + const Value * freq_factors = nullptr) { + if (status == nullptr) { + return; + } + status->log("ROPE matcher rejected node: %s input=%s positions=%s output=%s freq_factors=%s", reason, + value_layout_for_log(input).c_str(), value_layout_for_log(positions).c_str(), + value_layout_for_log(output).c_str(), value_layout_for_log(freq_factors).c_str()); +} + +static void log_set_rows_reject(Status * status, + const char * reason, + const Value * rows, + const Value * indices, + const Value * cache, + const Value * output) { + if (status == nullptr) { + return; + } + status->log("SET_ROWS matcher rejected node: %s rows=%s indices=%s cache=%s output=%s", reason, + value_layout_for_log(rows).c_str(), value_layout_for_log(indices).c_str(), + value_layout_for_log(cache).c_str(), value_layout_for_log(output).c_str()); +} + +static bool is_rope_shape(const Value & value) { + return value.ne[0] >= 4 && value.ne[0] <= 1024 && value.ne[0] % 4 == 0 && value.ne[1] >= 1 && value.ne[1] <= 64 && + value.ne[2] >= 1 && value.ne[2] <= 2048 && value.ne[3] == 1; +} + +static bool is_packed_f32_rope_layout(const Value & value) { + if (value.type != GGML_TYPE_F32 || !is_rope_shape(value) || value.nb[0] != sizeof(float)) { + return false; + } + return value.nb[1] == static_cast(value.ne[0]) * sizeof(float) && + value.nb[2] == static_cast(value.ne[1]) * value.nb[1] && + value.nb[3] == static_cast(value.ne[2]) * value.nb[2]; +} + +static bool is_supported_rope_input_layout(const Value & value, size_t & span_elements) { + if (!is_rope_shape(value) || value.nb[1] % sizeof(float) != 0 || value.nb[2] % sizeof(float) != 0 || + !strided_f32_storage_span_elements(value, span_elements)) { + return false; + } + + const size_t stride1 = value.nb[1] / sizeof(float); + const size_t stride2 = value.nb[2] / sizeof(float); + if (stride1 < static_cast(value.ne[0]) || stride2 < static_cast(value.ne[1]) * stride1) { + return false; + } + + const size_t token_span = static_cast(value.ne[2] - 1) * stride2; + const size_t head_span = static_cast(value.ne[1] - 1) * stride1; + const size_t required = token_span + head_span + static_cast(value.ne[0]); + return required <= span_elements; +} + +static bool is_supported_rope_params(const RopeParams & params, int64_t head_size) { + const bool supported_mode = params.mode == GGML_ROPE_TYPE_NORMAL || params.mode == GGML_ROPE_TYPE_NEOX; + return params.n_dims >= 4 && params.n_dims <= head_size && params.n_dims % 4 == 0 && supported_mode && + std::isfinite(params.freq_base) && params.freq_base > 0.0f && std::isfinite(params.freq_scale) && + params.freq_scale > 0.0f && std::isfinite(params.ext_factor) && std::isfinite(params.attn_factor); +} + +static bool build_rope_theta_table(const GraphNode & rope, + int64_t n_dims, + std::vector & data, + float & mscale) { + const RopeParams * params = op_params_as(rope.params); + if (params == nullptr) { + return false; + } + + RopeFrequencyTable table; + if (!build_rope_frequency_table(*params, n_dims, table)) { + return false; + } + data = std::move(table.data); + mscale = table.mscale; + return true; +} + +static void build_unit_frequency_factors(int64_t n_dims, std::vector & data) { + data.resize(static_cast(n_dims / 2) * sizeof(float)); + const float one = 1.0f; + for (int64_t i = 0; i < n_dims / 2; ++i) { + std::memcpy(data.data() + static_cast(i) * sizeof(float), &one, sizeof(one)); + } +} + +static bool format_value(ggml_type type, int64_t & value) { + switch (type) { + case GGML_TYPE_F16: + value = 16; + return true; + case GGML_TYPE_F32: + value = 32; + return true; + default: + return false; + } +} + +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; +} + +static bool supported_hidden_size(int64_t hidden_size) { + return hidden_size >= 4 && hidden_size <= 32768 && hidden_size % 4 == 0; +} + +static bool set_rows_supported_row_layout(const Value & rows, + int64_t hidden_size, + int64_t token_count, + int64_t & input_stride) { + if (!is_2d_shape(rows, hidden_size, token_count)) { + return false; + } + + const size_t element_size = ggml_type_size(rows.type); + if (rows.nb[0] != element_size || rows.nb[1] % element_size != 0) { + return false; + } + + input_stride = static_cast(rows.nb[1] / element_size); + return input_stride >= hidden_size && input_stride <= 1048576; +} + +static bool supported_set_rows_input_layout(const Value & rows, + int64_t hidden_size, + int64_t token_count, + int64_t & input_stride, + size_t & rows_span_bytes) { + if (!set_rows_supported_row_layout(rows, hidden_size, token_count, input_stride)) { + return false; + } + + if (rows.contiguous) { + rows_span_bytes = rows.byte_count; + return true; + } + + if (rows.type != GGML_TYPE_F32) { + return false; + } + + return strided_f32_storage_span_bytes(rows, rows_span_bytes); +} + +static ValueId next_match_transient_value(const DispatchMatchContext & context, const DispatchMatch & dispatch_match) { + return ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size()) + + static_cast(dispatch_match.completion_counter_requests.size())); +} + +struct RopeMatch { + const GraphNode * node = nullptr; + const Value * input = nullptr; + const Value * positions = nullptr; + const Value * freq_factors = nullptr; + const Value * output = nullptr; + int64_t token_count = 0; + int64_t head_count = 0; + int64_t head_size = 0; + int64_t n_dims = 0; + int64_t input_stride1 = 0; + int64_t input_stride2 = 0; + size_t input_span = 0; + size_t input_span_bytes = 0; + float rope_mscale = 1.0f; + int64_t mode = 0; + size_t theta_bytes = 0; + size_t freq_factors_bytes = 0; + std::vector theta_data; + std::vector freq_factors_data; + + bool matched() const { return node != nullptr && input != nullptr && positions != nullptr && output != nullptr; } +}; + +struct SetRowsMatch { + const GraphNode * node = nullptr; + const Value * rows = nullptr; + const Value * indices = nullptr; + const Value * cache = nullptr; + const Value * output = nullptr; + size_t rows_span_bytes = 0; + int64_t input_stride = 0; + int64_t row_format = 0; + int64_t output_format = 0; + int64_t token_count = 0; + int64_t cache_row_count = 0; + int64_t hidden_size = 0; + + bool matched() const { return node != nullptr && rows != nullptr && indices != nullptr && output != nullptr; } +}; + +struct RopeSetRowsMatch { + RopeMatch rope; + SetRowsMatch set_rows; + const GraphNode * layout = nullptr; + const Value * published_rows = nullptr; + size_t set_rows_idx = 0; + size_t layout_idx = 0; + + bool matched() const { return rope.matched() && set_rows.matched() && published_rows != nullptr; } +}; + +struct RopeScaleMatch { + RopeMatch rope; + const GraphNode * scale_node = nullptr; + const Value * scales = nullptr; + const Value * output = nullptr; + float scalar = 1.0f; + bool per_token = false; + + bool matched() const { return rope.matched() && scale_node != nullptr && output != nullptr; } +}; + +struct RopeConcatMatch { + RopeMatch rope; + const GraphNode * concat = nullptr; + const Value * prefix = nullptr; + const Value * output = nullptr; + size_t concat_idx = 0; + size_t input_span_bytes = 0; + + bool matched() const { return rope.matched() && concat != nullptr && prefix != nullptr && output != nullptr; } +}; + +static RopeMatch match_rope_f32(const Graph & graph, const GraphNode * node, Status * status = nullptr) { + RopeMatch match; + if (node == nullptr || node->op != GGML_OP_ROPE || node->inputs.size() < 2 || node->inputs.size() > 3) { + log_rope_reject(status, "root is not a supported ROPE arity", nullptr, nullptr, nullptr); + return match; + } + + const Value * input = graph_value(graph, node->inputs[0]); + const Value * positions = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + size_t input_span = 0; + if (input == nullptr || positions == nullptr || output == nullptr || input->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || positions->type != GGML_TYPE_I32 || + !is_supported_rope_input_layout(*input, input_span) || !is_packed_f32_rope_layout(*output) || + !positions->contiguous || !same_shape(*input, *output)) { + log_rope_reject(status, "input/output/positions shape or layout is unsupported", input, positions, output); + return {}; + } + + const int64_t head_size = input->ne[0]; + const int64_t head_count = input->ne[1]; + const int64_t token_count = input->ne[2]; + const RopeParams * params = op_params_as(node->params); + if (params == nullptr || !is_supported_rope_params(*params, head_size) || !is_1d_shape(*positions, token_count)) { + std::string reason = "ROPE params or positions shape is unsupported"; + if (params == nullptr) { + reason += " params=missing"; + } else { + std::ostringstream out; + out << reason << " params={n_dims=" << params->n_dims << ", mode=" << params->mode + << ", n_ctx_orig=" << params->n_ctx_orig << ", freq_base=" << params->freq_base + << ", freq_scale=" << params->freq_scale << ", ext_factor=" << params->ext_factor + << ", attn_factor=" << params->attn_factor << ", beta_fast=" << params->beta_fast + << ", beta_slow=" << params->beta_slow << "}"; + reason = out.str(); + } + log_rope_reject(status, reason.c_str(), input, positions, output); + return {}; + } + const int64_t n_dims = params->n_dims; + + std::vector theta_data; + float rope_mscale = 1.0f; + if (!build_rope_theta_table(*node, n_dims, theta_data, rope_mscale)) { + log_rope_reject(status, "failed to build supported ROPE theta table", input, positions, output); + return {}; + } + + const Value * freq_factors = nullptr; + size_t freq_factors_bytes = 0; + std::vector freq_factors_data; + if (node->inputs.size() == 3) { + freq_factors = graph_value(graph, node->inputs[2]); + if (freq_factors == nullptr || freq_factors->type != GGML_TYPE_F32 || !freq_factors->contiguous || + !is_1d_shape(*freq_factors, n_dims / 2)) { + log_rope_reject(status, "explicit frequency factors shape or layout is unsupported", input, positions, + output, freq_factors); + return {}; + } + freq_factors_bytes = freq_factors->byte_count; + } else { + build_unit_frequency_factors(n_dims, freq_factors_data); + freq_factors_bytes = freq_factors_data.size(); + } + + match.node = node; + match.input = input; + match.positions = positions; + match.freq_factors = freq_factors; + match.output = output; + match.token_count = token_count; + match.head_count = head_count; + match.head_size = head_size; + match.n_dims = n_dims; + match.input_stride1 = static_cast(input->nb[1] / sizeof(float)); + match.input_stride2 = static_cast(input->nb[2] / sizeof(float)); + match.input_span = input_span; + match.input_span_bytes = input_span * sizeof(float); + match.rope_mscale = rope_mscale; + match.mode = params->mode; + match.theta_bytes = theta_data.size(); + match.freq_factors_bytes = freq_factors_bytes; + match.theta_data = std::move(theta_data); + match.freq_factors_data = std::move(freq_factors_data); + return match; +} + +static SetRowsMatch match_set_rows_2d(const Graph & graph, const GraphNode * node, Status * status = nullptr) { + SetRowsMatch match; + if (node == nullptr || node->op != GGML_OP_SET_ROWS || node->inputs.size() != 3) { + log_set_rows_reject(status, "root is not a supported SET_ROWS arity", nullptr, nullptr, nullptr, nullptr); + return match; + } + + const Value * rows = graph_value(graph, node->inputs[0]); + const Value * indices = graph_value(graph, node->inputs[1]); + const Value * cache = graph_value(graph, node->inputs[2]); + const Value * output = graph_value(graph, node->output); + if (rows == nullptr || indices == nullptr || cache == nullptr || output == nullptr || !indices->contiguous || + !cache->contiguous || indices->type != GGML_TYPE_I64 || output->type != cache->type || + !same_shape(*cache, *output) || !graph.values().same_storage(cache->id, output->id)) { + log_set_rows_reject(status, "input/output shape or storage is unsupported", rows, indices, cache, output); + return {}; + } + + int64_t row_format = 0; + int64_t output_format = 0; + if (!format_value(rows->type, row_format) || !format_value(output->type, output_format)) { + log_set_rows_reject(status, "input or output format is unsupported", rows, indices, cache, output); + return {}; + } + if (rows->type == GGML_TYPE_F16 && output->type != GGML_TYPE_F16) { + log_set_rows_reject(status, "f16 input requires f16 output", rows, indices, cache, output); + return {}; + } + + const int64_t hidden_size = rows->ne[0]; + const int64_t token_count = rows->ne[1]; + const int64_t cache_row_count = cache->ne[1]; + int64_t input_stride = 0; + size_t rows_span_bytes = 0; + if (!supported_set_rows_input_layout(*rows, hidden_size, token_count, input_stride, rows_span_bytes) || + !is_1d_shape(*indices, token_count) || !is_2d_shape(*cache, hidden_size, cache_row_count) || + !supported_hidden_size(hidden_size) || !supported_token_count(token_count) || + !supported_cache_row_count(cache_row_count)) { + log_set_rows_reject(status, "shape or rows layout is unsupported", rows, indices, cache, output); + return {}; + } + + match.node = node; + match.rows = rows; + match.indices = indices; + match.cache = cache; + match.output = output; + match.rows_span_bytes = rows_span_bytes; + match.input_stride = input_stride; + match.row_format = row_format; + match.output_format = output_format; + match.token_count = token_count; + match.cache_row_count = cache_row_count; + match.hidden_size = hidden_size; + return match; +} + +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)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".token_capacity", + to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_stride1", + to_config_value(match.input_stride1)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + ".input_stride2", + to_config_value(match.input_stride2)); + 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)); +} + +static void add_set_rows_compile_parameters(Dispatch & dispatch, const SetRowsMatch & match) { + dispatch.kernel.compile_parameters.emplace("ggml.set_rows.token_capacity", to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.set_rows.hidden_capacity", to_config_value(match.hidden_size)); + dispatch.kernel.compile_parameters.emplace("ggml.set_rows.input_format", to_config_value(match.row_format)); + dispatch.kernel.compile_parameters.emplace("ggml.set_rows.output_format", to_config_value(match.output_format)); + dispatch.kernel.compile_parameters.emplace("ggml.set_rows.input_stride", to_config_value(match.input_stride)); +} + +static ValueId add_constant_binding(const DispatchMatchContext & context, + DispatchMatch & dispatch_match, + const char * name, + size_t byte_count, + const std::vector & data) { + const ValueId value = next_match_transient_value(context, dispatch_match); + dispatch_match.transients.push_back({ value, name, byte_count, 256 }); + dispatch_match.constant_initializations.push_back({ + value, + name, + 0, + data, + }); + return value; +} + +static std::pair add_rope_frequency_bindings(const DispatchMatchContext & context, + const RopeMatch & match, + DispatchMatch & dispatch_match, + const char * prefix) { + const std::string theta_name = std::string(prefix) + ".theta"; + const ValueId theta = + add_constant_binding(context, dispatch_match, theta_name.c_str(), match.theta_bytes, match.theta_data); + if (match.freq_factors != nullptr) { + return { theta, match.freq_factors->id }; + } + + const std::string factors_name = std::string(prefix) + ".freq_factors"; + const ValueId freq_factors = add_constant_binding(context, dispatch_match, factors_name.c_str(), + match.freq_factors_bytes, match.freq_factors_data); + return { theta, freq_factors }; +} + +static Dispatch make_rope_dispatch(const DispatchMatchContext & context, + const RopeMatch & match, + DispatchMatch & dispatch_match) { + 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); + 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 }); + dispatch.bindings.push_back({ freq_factors, 0, match.freq_factors_bytes }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + return dispatch; +} + +static Dispatch make_rope_scale_dispatch(const DispatchMatchContext & context, + const RopeScaleMatch & match, + DispatchMatch & dispatch_match) { + const auto [theta, freq_factors] = + add_rope_frequency_bindings(context, match.rope, dispatch_match, "common.rope_scale"); + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(match.per_token ? kRopeTokenScaleF32Kernel : kRopeScaleF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.rope.token_count); + add_rope_compile_parameters(dispatch, match.rope, "ggml.rope_f32"); + if (!match.per_token) { + dispatch.kernel.compile_parameters.emplace("ggml.rope_scale_f32.scale", rope_mscale_config_value(match.scalar)); + } + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back( + { match.rope.input->storage_root, match.rope.input->storage_offset, match.rope.input_span_bytes }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + if (match.per_token) { + dispatch.bindings.push_back( + { match.scales->storage_root, match.scales->storage_offset, match.scales->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + return dispatch; +} + +static Dispatch make_set_rows_dispatch(const SetRowsMatch & match) { + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kSetRowsKernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); + dispatch.kernel.integer_parameters.emplace("cache_row_count", match.cache_row_count); + dispatch.kernel.integer_parameters.emplace("hidden_size", match.hidden_size); + add_set_rows_compile_parameters(dispatch, match); + dispatch.bindings.push_back({ match.rows->storage_root, match.rows->storage_offset, match.rows_span_bytes }); + dispatch.bindings.push_back({ match.indices->id, 0, match.indices->byte_count }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + return dispatch; +} + +static Dispatch make_rope_set_rows_dispatch(const DispatchMatchContext & context, + const RopeSetRowsMatch & match, + DispatchMatch & dispatch_match) { + const auto [theta, freq_factors] = + add_rope_frequency_bindings(context, match.rope, dispatch_match, "common.rope_set_rows"); + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRopeSetRowsF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.rope.token_count); + dispatch.kernel.integer_parameters.emplace("cache_row_count", match.set_rows.cache_row_count); + add_rope_compile_parameters(dispatch, match.rope, "ggml.rope_set_rows_f32"); + dispatch.kernel.compile_parameters.emplace("ggml.rope_set_rows_f32.output_format", + to_config_value(match.set_rows.output_format)); + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back({ match.set_rows.indices->id, 0, match.set_rows.indices->byte_count }); + dispatch.bindings.push_back( + { match.rope.input->storage_root, match.rope.input->storage_offset, match.rope.input_span_bytes }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + dispatch.bindings.push_back({ match.set_rows.output->id, 0, match.set_rows.output->byte_count }); + return dispatch; +} + +static Dispatch make_rope_concat_dispatch(const DispatchMatchContext & context, + const RopeConcatMatch & match, + DispatchMatch & dispatch_match) { + const auto [theta, freq_factors] = + add_rope_frequency_bindings(context, match.rope, dispatch_match, "common.rope_concat"); + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kRopeConcatF32Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.rope.token_count); + const char * prefix = "ggml.rope_concat_f32."; + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "prefix_size", + to_config_value(match.prefix->ne[0])); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "rope_size", + to_config_value(match.rope.head_size)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "n_dims", to_config_value(match.rope.n_dims)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "head_count", + to_config_value(match.rope.head_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "token_capacity", + to_config_value(match.rope.token_count)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "input_stride1", + to_config_value(match.rope.input_stride1)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "input_stride2", + to_config_value(match.rope.input_stride2)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "input_span", + to_config_value(match.input_span_bytes / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "mscale", + rope_mscale_config_value(match.rope.rope_mscale)); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "mode", to_config_value(match.rope.mode)); + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back({ match.prefix->storage_root, match.prefix->storage_offset, match.input_span_bytes }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + return dispatch; +} + +static RopeConcatMatch match_rope_concat_f32(const Graph & graph, const GraphNode * node) { + RopeConcatMatch match; + match.rope = match_rope_f32(graph, node); + if (!match.rope.matched() || match.rope.mode != GGML_ROPE_TYPE_NEOX || !graph.has_index() || + !graph.index().has_single_consumer(node->output)) { + return {}; + } + + const GraphNode * concat = graph.index().consumers(node->output).front(); + if (concat == nullptr || concat->op != GGML_OP_CONCAT || concat->inputs.size() != 2 || + concat->inputs[1] != node->output || !graph.index().node_index(concat, match.concat_idx)) { + return {}; + } + + const Value * prefix = graph_value(graph, concat->inputs[0]); + const Value * output = graph_value(graph, concat->output); + if (prefix == nullptr || output == nullptr || prefix->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + prefix->ne[0] < 4 || prefix->ne[0] > 1024 || prefix->ne[0] % 4 != 0 || + prefix->ne[0] + match.rope.head_size > 1024 || !is_packed_f32_rope_layout(*output) || + output->ne[0] != prefix->ne[0] + match.rope.head_size || prefix->nb[0] != sizeof(float) || + prefix->storage_root != match.rope.input->storage_root || prefix->storage_root == output->storage_root || + match.rope.input->storage_offset != + prefix->storage_offset + static_cast(prefix->ne[0]) * sizeof(float)) { + return {}; + } + for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) { + if (prefix->ne[dim] != match.rope.input->ne[dim] || prefix->ne[dim] != output->ne[dim] || + prefix->nb[dim] != match.rope.input->nb[dim]) { + return {}; + } + } + + size_t prefix_span_bytes = 0; + if (!strided_f32_storage_span_bytes(*prefix, prefix_span_bytes)) { + return {}; + } + const size_t rope_offset_bytes = match.rope.input->storage_offset - prefix->storage_offset; + match.input_span_bytes = std::max(prefix_span_bytes, rope_offset_bytes + match.rope.input_span_bytes); + match.concat = concat; + match.prefix = prefix; + match.output = output; + return match; +} + +static bool match_rope_concat_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RopeConcatMatch match = match_rope_concat_f32(context.graph, context.root_node); + if (!match.matched() || + !append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.concat, + dispatch_match.covered_nodes)) { + return false; + } + + dispatch_match.dispatches.push_back(make_rope_concat_dispatch(context, match, dispatch_match)); + return true; +} + +static bool find_rope_set_rows_consumer(const Graph & graph, const GraphNode & rope, RopeSetRowsMatch & match) { + if (!graph.has_index() || !graph.index().has_single_consumer(rope.output)) { + return false; + } + + const GraphNode * consumer = graph.index().consumers(rope.output).front(); + if (consumer == nullptr) { + return false; + } + + if (consumer->op == GGML_OP_SET_ROWS) { + match.set_rows_idx = 0; + match.layout = nullptr; + return graph.index().node_index(consumer, match.set_rows_idx) && + (match.set_rows = match_set_rows_2d(graph, consumer)).matched(); + } + + if (!is_layout_alias_node(graph, *consumer) || !graph.index().has_single_consumer(consumer->output)) { + return false; + } + + const GraphNode * set_rows = graph.index().consumers(consumer->output).front(); + if (set_rows == nullptr || set_rows->op != GGML_OP_SET_ROWS) { + return false; + } + match.layout = consumer; + return graph.index().node_index(consumer, match.layout_idx) && + graph.index().node_index(set_rows, match.set_rows_idx) && + (match.set_rows = match_set_rows_2d(graph, set_rows)).matched(); +} + +static RopeSetRowsMatch match_rope_set_rows_f32(const Graph & graph, const GraphNode * node) { + RopeSetRowsMatch match; + match.rope = match_rope_f32(graph, node); + if (!match.rope.matched() || !find_rope_set_rows_consumer(graph, *node, match)) { + return {}; + } + + match.published_rows = match.set_rows.rows; + if (match.set_rows.row_format != 32 || match.set_rows.hidden_size != match.rope.head_size * match.rope.head_count || + match.set_rows.token_count != match.rope.token_count || match.published_rows->type != GGML_TYPE_F32) { + return {}; + } + return match; +} + +static RopeScaleMatch match_rope_scale_f32(const Graph & graph, const GraphNode * node) { + RopeScaleMatch match; + match.rope = match_rope_f32(graph, node); + if (!match.rope.matched() || !graph.has_index() || !graph.index().has_single_consumer(node->output)) { + return {}; + } + + const GraphNode * scale_node = graph.index().consumers(node->output).front(); + if (scale_node == nullptr) { + return {}; + } + const Value * output = graph_value(graph, scale_node->output); + if (output == nullptr || output->type != GGML_TYPE_F32 || !output->contiguous || output->alias_source.value >= 0 || + !same_shape(*match.rope.output, *output)) { + return {}; + } + + if (scale_node->op == GGML_OP_SCALE && scale_node->inputs.size() == 1 && + scale_node->inputs.front() == node->output) { + const ScaleParams * params = op_params_as(scale_node->params); + if (params == nullptr || params->bias != 0.0f || !std::isfinite(params->scale)) { + return {}; + } + match.scalar = params->scale; + } else if (scale_node->op == GGML_OP_MUL && scale_node->inputs.size() == 2) { + const BinaryParams * params = op_params_as(scale_node->params); + if (params == nullptr || params->op != BinaryKind::Mul) { + return {}; + } + ValueId scale_id; + if (scale_node->inputs[0] == node->output) { + scale_id = scale_node->inputs[1]; + } else if (scale_node->inputs[1] == node->output) { + scale_id = scale_node->inputs[0]; + } else { + return {}; + } + const Value * scales = graph_value(graph, scale_id); + if (scales == nullptr || scales->type != GGML_TYPE_F32 || !scales->contiguous || scales->ne[0] != 1 || + scales->ne[1] != 1 || scales->ne[2] != match.rope.token_count || scales->ne[3] != 1) { + return {}; + } + match.scales = scales; + match.per_token = true; + } else { + return {}; + } + + match.scale_node = scale_node; + match.output = output; + return match; +} + +static bool value_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +static bool match_rope_scale_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RopeScaleMatch match = match_rope_scale_f32(context.graph, context.root_node); + if (!match.matched() || + (match.per_token && !value_is_available(context.graph, match.scales->id, context.covered_nodes)) || + !append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.scale_node, + dispatch_match.covered_nodes)) { + return false; + } + + dispatch_match.dispatches.push_back(make_rope_scale_dispatch(context, match, dispatch_match)); + return true; +} + +static bool match_rope_set_rows_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RopeSetRowsMatch match = match_rope_set_rows_f32(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes) || + (match.layout != nullptr && !append_covered_node_index_once(context.graph, context.covered_nodes, match.layout, + dispatch_match.covered_nodes)) || + !append_covered_node_index_once(context.graph, context.covered_nodes, match.set_rows.node, + dispatch_match.covered_nodes)) { + return false; + } + + dispatch_match.dispatches.push_back(make_rope_set_rows_dispatch(context, match, dispatch_match)); + return true; +} + +static bool match_rope_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const RopeMatch match = match_rope_f32(context.graph, context.root_node, &dispatch_match.status); + if (!match.matched()) { + return false; + } + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(make_rope_dispatch(context, match, dispatch_match)); + return true; +} + +static bool match_set_rows_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const SetRowsMatch match = match_set_rows_2d(context.graph, context.root_node, &dispatch_match.status); + if (!match.matched()) { + return false; + } + + dispatch_match.covered_nodes.push_back(context.root_index); + dispatch_match.dispatches.push_back(make_set_rows_dispatch(match)); + return true; +} + +} // namespace + +void register_rope_set_rows_dispatches(DispatchRegistryBuilder & registry) { + registry.add({ + "common.rope_scale.f32", + GGML_OP_ROPE, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_rope_scale_f32_dispatch, + }); + registry.add({ + "common.rope_concat.f32", + GGML_OP_ROPE, + DispatchMatchKind::Fused, + 300, + DispatchSource::Common, + match_rope_concat_f32_dispatch, + }); + registry.add({ + "common.rope_set_rows.f32", + GGML_OP_ROPE, + DispatchMatchKind::Fused, + 200, + DispatchSource::Common, + match_rope_set_rows_f32_dispatch, + }); + registry.add({ + "common.rope.f32", + GGML_OP_ROPE, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Common, + match_rope_f32_dispatch, + }); + registry.add({ + "common.set_rows", + GGML_OP_SET_ROWS, + DispatchMatchKind::SingleOp, + 100, + DispatchSource::Common, + match_set_rows_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.h new file mode 100644 index 000000000000..8e9497a6fe68 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-set-rows.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_rope_set_rows_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-utils.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-utils.h new file mode 100644 index 000000000000..7981ccd1c03e --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-rope-utils.h @@ -0,0 +1,72 @@ +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { + +struct RopeFrequencyTable { + std::vector data; + float mscale = 1.0f; +}; + +inline float rope_yarn_ramp(float low, float high, int64_t i0) { + const float y = (static_cast(i0 / 2) - low) / std::max(0.001f, high - low); + return 1.0f - std::min(1.0f, std::max(0.0f, y)); +} + +inline float rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { + return static_cast(n_dims) * std::log(static_cast(n_ctx_orig) / (n_rot * 2.0f * static_cast(M_PI))) / + (2.0f * std::log(base)); +} + +inline void rope_yarn_corr_dims(int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]) { + const float start = std::floor(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); + const float end = std::ceil(rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); + dims[0] = std::max(0.0f, start); + dims[1] = std::min(static_cast(n_dims - 1), end); +} + +inline bool build_rope_frequency_table(const RopeParams & params, int64_t n_dims, RopeFrequencyTable & table) { + if (n_dims <= 0 || n_dims % 2 != 0 || params.freq_base <= 0.0f || !std::isfinite(params.freq_base) || + params.freq_scale <= 0.0f || !std::isfinite(params.freq_scale) || !std::isfinite(params.ext_factor) || + !std::isfinite(params.attn_factor)) { + return false; + } + + float corr_dims[2] = { 0.0f, 0.0f }; + rope_yarn_corr_dims(static_cast(n_dims), params.n_ctx_orig, params.freq_base, params.beta_fast, params.beta_slow, + corr_dims); + + table.mscale = params.attn_factor; + if (params.ext_factor != 0.0f) { + table.mscale *= 1.0f + 0.1f * std::log(1.0f / params.freq_scale); + } + + table.data.resize(static_cast(n_dims / 2) * sizeof(float)); + const float theta_scale = std::pow(params.freq_base, -2.0f / static_cast(n_dims)); + float theta = 1.0f; + for (int64_t i = 0; i < n_dims / 2; ++i) { + const int64_t dim = 2 * i; + const float ramp_mix = rope_yarn_ramp(corr_dims[0], corr_dims[1], dim) * params.ext_factor; + const float frequency = params.freq_scale * theta * (1.0f - ramp_mix) + theta * ramp_mix; + std::memcpy(table.data.data() + static_cast(i) * sizeof(float), &frequency, sizeof(frequency)); + theta *= theta_scale; + } + return true; +} + +inline std::string rope_mscale_config_value(float value) { + char buffer[64]; + std::snprintf(buffer, sizeof(buffer), "%.9g", static_cast(value)); + return std::string(buffer); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp new file mode 100644 index 000000000000..e7cd7f67f91c --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.cpp @@ -0,0 +1,256 @@ +#include "dispatch-scale.h" + +#include "dispatch-binary-common.h" +#include "dispatch-layout-utils.h" +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kScaleF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_f32"); +static constexpr KernelCatalogRef kScaleAddF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_add_f32"); + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static bool positive_shape(const Value & value) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.ne[i] <= 0) { + return false; + } + } + return value.element_count > 0; +} + +static bool packed_f32_layout(const Value & value) { + size_t expected_stride = sizeof(float); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (value.nb[i] != expected_stride) { + return false; + } + expected_stride *= static_cast(value.ne[i]); + } + return true; +} + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool distinct_storage(const Value & input, const Value & output) { + return input.storage != output.storage; +} + +static bool storage_ranges_disjoint(const Value & lhs, + size_t lhs_byte_count, + const Value & rhs, + size_t rhs_byte_count) { + if (lhs.storage != rhs.storage) { + return true; + } + if (lhs.storage_offset > std::numeric_limits::max() - lhs_byte_count || + rhs.storage_offset > std::numeric_limits::max() - rhs_byte_count) { + return false; + } + return lhs.storage_offset + lhs_byte_count <= rhs.storage_offset || + rhs.storage_offset + rhs_byte_count <= lhs.storage_offset; +} + +static bool value_is_available(const Graph & graph, ValueId value, const std::vector & covered_nodes) { + const GraphNode * producer = graph.index().producer(value); + if (producer == nullptr) { + return true; + } + size_t producer_index = 0; + return graph.index().node_index(producer, producer_index) && producer_index < covered_nodes.size() && + covered_nodes[producer_index]; +} + +static bool value_is_available_or_layout_alias(const Graph & graph, + ValueId value, + const std::vector & covered_nodes) { + if (value_is_available(graph, value, covered_nodes)) { + return true; + } + const GraphNode * producer = graph.index().producer(value); + return producer != nullptr && is_layout_alias_node(graph, *producer) && producer->inputs.size() == 1 && + value_is_available(graph, producer->inputs[0], covered_nodes); +} + +static bool supported_source_layout(const Graph & graph, const Value & value) { + if (value.alias_source.value < 0) { + return true; + } + if (value.storage_offset != 0) { + return false; + } + const GraphNode * producer = graph.index().producer(value.id); + return producer != nullptr && (producer->op == GGML_OP_RESHAPE || producer->op == GGML_OP_CLAMP); +} + +static std::string format_float_config(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + +static void add_strided_source_config(Dispatch & dispatch, + const char * prefix, + const Value & source, + size_t source_byte_count) { + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "stride1", + std::to_string(source.nb[1] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "stride2", + std::to_string(source.nb[2] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "stride3", + std::to_string(source.nb[3] / sizeof(float))); + dispatch.kernel.compile_parameters.emplace(std::string(prefix) + "span", + std::to_string(source_byte_count / sizeof(float))); +} + +static bool match_scale_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * scale_node = context.root_node; + if (scale_node == nullptr || scale_node->op != GGML_OP_SCALE || scale_node->inputs.size() != 1 || + !context.graph.has_index() || context.root_index >= context.covered_nodes.size() || + context.covered_nodes[context.root_index]) { + return false; + } + + const ScaleParams * scale_params = op_params_as(scale_node->params); + if (scale_params == nullptr || !context.graph.index().has_single_consumer(scale_node->output)) { + return false; + } + const GraphNode * binary_node = context.graph.index().consumers(scale_node->output).front(); + const CommonBinaryMatch binary = + common_match_binary_consumer(binary_node, scale_node->output, kCommonArithmeticBinaryKinds); + if (!binary.matched()) { + return false; + } + + size_t binary_index = 0; + if (!context.graph.index().node_index(binary_node, binary_index) || binary_index >= context.covered_nodes.size() || + context.covered_nodes[binary_index]) { + return false; + } + + const Value * input = graph_value(context.graph, scale_node->inputs[0]); + const Value * scaled = graph_value(context.graph, scale_node->output); + const Value * residual = graph_value(context.graph, binary.operand); + const Value * output = graph_value(context.graph, binary_node->output); + if (input == nullptr || scaled == nullptr || residual == nullptr || output == nullptr || + input->type != GGML_TYPE_F32 || scaled->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || + output->type != GGML_TYPE_F32 || !same_shape(*input, *scaled) || !same_shape(*input, *residual) || + !same_shape(*input, *output) || !positive_shape(*output) || !output->contiguous || + !packed_f32_layout(*output) || scaled->alias_source.value >= 0 || output->alias_source.value >= 0 || + static_cast(output->element_count) > std::numeric_limits::max()) { + return false; + } + + size_t input_byte_count = 0; + size_t residual_byte_count = 0; + if (!strided_f32_storage_span_bytes(*input, input_byte_count) || + !strided_f32_storage_span_bytes(*residual, residual_byte_count) || + !storage_ranges_disjoint(*input, input_byte_count, *output, output->byte_count) || + !storage_ranges_disjoint(*residual, residual_byte_count, *output, output->byte_count) || + !value_is_available_or_layout_alias(context.graph, input->id, context.covered_nodes) || + !value_is_available_or_layout_alias(context.graph, residual->id, context.covered_nodes)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kScaleAddF32Kernel); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.ne0", std::to_string(output->ne[0])); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.ne1", std::to_string(output->ne[1])); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.ne2", std::to_string(output->ne[2])); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.scale", format_float_config(scale_params->scale)); + dispatch.kernel.compile_parameters.emplace("ggml.scale_add_f32.bias", format_float_config(scale_params->bias)); + common_set_binary_compile_parameters(dispatch.kernel, "ggml.scale_add_f32.binary_op", + "ggml.scale_add_f32.scaled_lhs", binary.kind, binary.producer_is_lhs); + add_strided_source_config(dispatch, "ggml.scale_add_f32.input_", *input, input_byte_count); + add_strided_source_config(dispatch, "ggml.scale_add_f32.residual_", *residual, residual_byte_count); + dispatch.bindings.push_back({ input->storage_root, input->storage_offset, input_byte_count }); + dispatch.bindings.push_back({ residual->storage_root, residual->storage_offset, residual_byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.covered_nodes.push_back(binary_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_scale_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->op != GGML_OP_SCALE || node->inputs.size() != 1) { + return false; + } + + const ScaleParams * params = op_params_as(node->params); + if (params == nullptr) { + return false; + } + + const Value * output = graph_value(context.graph, node->output); + const Value * input = graph_value(context.graph, node->inputs[0]); + if (output == nullptr || input == nullptr) { + return false; + } + + if (output->type != GGML_TYPE_F32 || input->type != GGML_TYPE_F32 || !same_shape(*output, *input) || + !positive_shape(*output) || !output->contiguous || !input->contiguous || !packed_f32_layout(*output) || + !packed_f32_layout(*input) || output->alias_source.value >= 0 || + !supported_source_layout(context.graph, *input) || !distinct_storage(*input, *output) || + static_cast(output->element_count) > std::numeric_limits::max()) { + return false; + } + + 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)); + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +} // namespace + +void register_scale_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.scale_add_f32", + GGML_OP_SCALE, + DispatchMatchKind::Fused, + 1000, + DispatchSource::Common, + match_scale_add_f32_dispatch, + }); + registry.add({ + "common.scale_f32", + GGML_OP_SCALE, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_scale_f32_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.h new file mode 100644 index 000000000000..7c962a0062fd --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-scale.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch_registration/dispatch-registry.h" + +namespace ggml::hrx { + +void register_scale_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..03cbc2b8b232 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.cpp @@ -0,0 +1,151 @@ +#include "dispatch-unary.h" + +#include "ggml.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kUnaryF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_unary_f32"); +static constexpr KernelCatalogRef kScaleBiasF32Kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_scale_bias_f32"); + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static std::string float_config_value(float value) { + std::ostringstream stream; + stream << std::setprecision(std::numeric_limits::max_digits10) << value; + return stream.str(); +} + +static bool scale_storage_is_safe(const Value & input, const Value & output) { + if (input.storage_root != output.storage_root) { + return true; + } + if (input.storage_offset == output.storage_offset && input.byte_count == output.byte_count) { + return true; + } + if (input.storage_offset > std::numeric_limits::max() - input.byte_count || + output.storage_offset > std::numeric_limits::max() - output.byte_count) { + return false; + } + return input.storage_offset + input.byte_count <= output.storage_offset || + output.storage_offset + output.byte_count <= input.storage_offset; +} + +static bool match_scale_bias_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->inputs.size() != 1) { + return false; + } + + const ScaleParams * params = op_params_as(node->params); + const Value * input = graph_value(context.graph, node->inputs[0]); + const Value * output = graph_value(context.graph, node->output); + if (params == nullptr || input == nullptr || output == nullptr || !std::isfinite(params->scale) || + !std::isfinite(params->bias) || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + !input->contiguous || !output->contiguous || !same_shape(*input, *output) || + !scale_storage_is_safe(*input, *output) || output->element_count <= 0 || + static_cast(output->element_count) > std::numeric_limits::max()) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kScaleBiasF32Kernel); + dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); + dispatch.kernel.compile_parameters.emplace("ggml.scale.scale", float_config_value(params->scale)); + dispatch.kernel.compile_parameters.emplace("ggml.scale.bias", float_config_value(params->bias)); + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_unary_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { + const GraphNode * node = context.root_node; + if (node == nullptr || node->inputs.size() != 1) { + return false; + } + + const UnaryParams * params = op_params_as(node->params); + if (params == nullptr || !unary_kind_supported(params->op)) { + return false; + } + + const Value * output = graph_value(context.graph, node->output); + const Value * input = graph_value(context.graph, node->inputs[0]); + if (output == nullptr || input == nullptr) { + return false; + } + + if (output->type != GGML_TYPE_F32 || input->type != GGML_TYPE_F32 || !same_shape(*output, *input) || + !output->contiguous || !input->contiguous || output->element_count <= 0 || + static_cast(output->element_count) > std::numeric_limits::max()) { + return false; + } + + 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))); + dispatch.bindings.push_back({ input->id, 0, input->byte_count }); + dispatch.bindings.push_back({ output->id, 0, output->byte_count }); + + match.covered_nodes.push_back(context.root_index); + match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static void register_unary_dispatch_for(DispatchRegistryBuilder & registry, ggml_op root_op) { + registry.add({ + "common.unary_f32", + root_op, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_unary_f32_dispatch, + }); +} + +} // namespace + +void register_unary_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "common.scale_bias_f32", + GGML_OP_SCALE, + DispatchMatchKind::SingleOp, + 0, + DispatchSource::Common, + match_scale_bias_f32_dispatch, + }); + register_unary_dispatch_for(registry, GGML_OP_UNARY); + register_unary_dispatch_for(registry, GGML_OP_SQR); + register_unary_dispatch_for(registry, GGML_OP_SQRT); + register_unary_dispatch_for(registry, GGML_OP_LOG); + register_unary_dispatch_for(registry, GGML_OP_SIN); + register_unary_dispatch_for(registry, GGML_OP_COS); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.h b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.h new file mode 100644 index 000000000000..c42d89d69f0b --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-unary.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_unary_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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-add.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.cpp deleted file mode 100644 index 49ecb5664578..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.cpp +++ /dev/null @@ -1,77 +0,0 @@ -#include "dispatch-add.h" - -#include "ggml.h" -#include "kernel-corpus/kernel-corpus-catalog-verify.h" - -#include -#include -#include - -namespace ggml::hrx { - -static constexpr KernelCatalogRef kAddF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_add_f32"); - -static bool same_shape(const Value & lhs, const Value & rhs) { - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (lhs.ne[i] != rhs.ne[i]) { - return false; - } - } - return true; -} - -static const Value * graph_value(const Graph & graph, ValueId id) { - return graph.values().find(id); -} - -static bool supports_add_f32_dispatch(const Graph & graph, const GraphNode * node) { - if (node == nullptr || node->op != GGML_OP_ADD || node->inputs.size() != 2) { - return false; - } - const Value * output = graph_value(graph, node->output); - const Value * a = graph_value(graph, node->inputs[0]); - const Value * b = graph_value(graph, node->inputs[1]); - if (output == nullptr || a == nullptr || b == nullptr) { - return false; - } - return output->type == GGML_TYPE_F32 && a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32 && - same_shape(*output, *a) && same_shape(*output, *b) && output->contiguous && a->contiguous && b->contiguous && - output->element_count > 0 && - static_cast(output->element_count) <= std::numeric_limits::max(); -} - -static bool match_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { - if (!supports_add_f32_dispatch(context.graph, context.root_node)) { - return false; - } - const Value * output = graph_value(context.graph, context.root_node->output); - const Value * a = graph_value(context.graph, context.root_node->inputs[0]); - const Value * b = graph_value(context.graph, context.root_node->inputs[1]); - if (output == nullptr || a == nullptr || b == nullptr) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kAddF32Kernel); - dispatch.kernel.integer_parameters.emplace("element_count", output->element_count); - dispatch.bindings.push_back({ a->id, 0, a->byte_count }); - dispatch.bindings.push_back({ b->id, 0, b->byte_count }); - dispatch.bindings.push_back({ output->id, 0, output->byte_count }); - - match.covered_nodes.push_back(context.root_index); - match.dispatches.push_back(std::move(dispatch)); - return true; -} - -void register_add_dispatch(DispatchRegistryBuilder & registry) { - registry.add({ - "common.add_f32", - GGML_OP_ADD, - DispatchMatchKind::SingleOp, - 0, - DispatchSource::Common, - match_add_f32_dispatch, - }); -} - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.h deleted file mode 100644 index 3054b942958c..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-add.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "dispatch-registry.h" - -namespace ggml::hrx { - -void register_add_dispatch(DispatchRegistryBuilder & registry); - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.cpp deleted file mode 100644 index 2bcd8b4d4616..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-gather-add.cpp +++ /dev/null @@ -1,200 +0,0 @@ -#include "dispatch-gather-add.h" - -#include "ggml.h" -#include "kernel-corpus/kernel-corpus-catalog-verify.h" - -#include -#include - -namespace ggml::hrx { -namespace { - -static constexpr KernelCatalogRef kGatherAddF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "ggml_gather_add_f32"); - -static const Value * graph_value(const Graph & graph, ValueId id) { - return graph.values().find(id); -} - -static bool same_shape(const Value & lhs, const Value & rhs) { - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (lhs.ne[i] != rhs.ne[i]) { - return false; - } - } - return true; -} - -static bool is_2d_f32(const Value & value) { - return value.type == GGML_TYPE_F32 && value.contiguous && value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && - value.ne[3] == 1; -} - -static bool is_row_id_tensor(const Value & value, int64_t output_token_count) { - return value.type == GGML_TYPE_I32 && value.contiguous && value.element_count == output_token_count && - output_token_count > 0; -} - -static bool is_supported_hidden_size(int64_t hidden_size) { - return hidden_size >= 128 && hidden_size <= 32768 && hidden_size % 128 == 0; -} - -static bool is_supported_token_count(int64_t token_count) { - return token_count >= 1 && token_count <= 2048; -} - -struct GatherAddMatch { - const GraphNode * first_get_rows = nullptr; - const GraphNode * second_get_rows = nullptr; - const GraphNode * add_node = nullptr; - const Value * first_source = nullptr; - const Value * second_source = nullptr; - const Value * row_ids = nullptr; - const Value * output = nullptr; - size_t first_get_rows_index = 0; - size_t second_get_rows_index = 0; - size_t add_node_index = 0; - int64_t source_token_count = 0; - int64_t output_token_count = 0; - int64_t hidden_size = 0; - - bool matched() const { - return first_get_rows != nullptr && second_get_rows != nullptr && add_node != nullptr && - first_source != nullptr && second_source != nullptr && row_ids != nullptr && output != nullptr; - } -}; - -static const GraphNode * find_single_add_consumer(const Graph & graph, const GraphNode & get_rows) { - const std::vector & consumers = graph.index().consumers(get_rows.output); - if (consumers.size() != 1) { - return nullptr; - } - const GraphNode * add = consumers.front(); - return add != nullptr && add->op == GGML_OP_ADD && add->inputs.size() == 2 ? add : nullptr; -} - -static const GraphNode * peer_get_rows_input(const Graph & graph, - const GraphNode & add, - const GraphNode & root_get_rows) { - ValueId peer_output; - if (add.inputs[0] == root_get_rows.output) { - peer_output = add.inputs[1]; - } else if (add.inputs[1] == root_get_rows.output) { - peer_output = add.inputs[0]; - } else { - return nullptr; - } - - const GraphNode * peer = graph.index().producer(peer_output); - return peer != nullptr && peer->op == GGML_OP_GET_ROWS && peer->inputs.size() == 2 ? peer : nullptr; -} - -static GatherAddMatch match_gather_add_f32(const Graph & graph, const GraphNode * node, size_t node_index) { - GatherAddMatch match; - if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2 || !graph.has_index()) { - return match; - } - - const GraphNode * add_node = find_single_add_consumer(graph, *node); - const GraphNode * peer = add_node != nullptr ? peer_get_rows_input(graph, *add_node, *node) : nullptr; - if (add_node == nullptr || peer == nullptr) { - return {}; - } - - size_t peer_index = 0; - size_t add_index = 0; - if (!graph.index().node_index(peer, peer_index) || !graph.index().node_index(add_node, add_index)) { - return {}; - } - - const Value * first_source = graph_value(graph, node->inputs[0]); - const Value * first_ids = graph_value(graph, node->inputs[1]); - const Value * first_output = graph_value(graph, node->output); - const Value * second_source = graph_value(graph, peer->inputs[0]); - const Value * second_ids = graph_value(graph, peer->inputs[1]); - const Value * second_output = graph_value(graph, peer->output); - const Value * output = graph_value(graph, add_node->output); - if (first_source == nullptr || first_ids == nullptr || first_output == nullptr || second_source == nullptr || - second_ids == nullptr || second_output == nullptr || output == nullptr) { - return {}; - } - if (node->inputs[1] != peer->inputs[1]) { - return {}; - } - if (!is_2d_f32(*first_source) || !is_2d_f32(*second_source) || !is_2d_f32(*first_output) || - !is_2d_f32(*second_output) || !is_2d_f32(*output)) { - return {}; - } - if (!same_shape(*first_source, *second_source) || !same_shape(*first_output, *second_output) || - !same_shape(*first_output, *output)) { - return {}; - } - - const int64_t hidden_size = first_source->ne[0]; - const int64_t source_token_count = first_source->ne[1]; - const int64_t output_token_count = first_output->ne[1]; - if (first_output->ne[0] != hidden_size || !is_row_id_tensor(*first_ids, output_token_count) || - !is_row_id_tensor(*second_ids, output_token_count)) { - return {}; - } - if (!is_supported_hidden_size(hidden_size) || !is_supported_token_count(source_token_count) || - !is_supported_token_count(output_token_count)) { - return {}; - } - - match.first_get_rows = node; - match.second_get_rows = peer; - match.add_node = add_node; - match.first_source = first_source; - match.second_source = second_source; - match.row_ids = first_ids; - match.output = output; - match.first_get_rows_index = node_index; - match.second_get_rows_index = peer_index; - match.add_node_index = add_index; - match.source_token_count = source_token_count; - match.output_token_count = output_token_count; - match.hidden_size = hidden_size; - return match; -} - -} // namespace - -static bool match_gather_add_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { - const GatherAddMatch gather_add = match_gather_add_f32(context.graph, context.root_node, context.root_index); - if (!gather_add.matched() || gather_add.first_get_rows_index >= context.covered_nodes.size() || - gather_add.second_get_rows_index >= context.covered_nodes.size() || - gather_add.add_node_index >= context.covered_nodes.size() || - context.covered_nodes[gather_add.first_get_rows_index] || - context.covered_nodes[gather_add.second_get_rows_index] || context.covered_nodes[gather_add.add_node_index]) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kGatherAddF32Kernel); - dispatch.kernel.integer_parameters.emplace("source_token_count", gather_add.source_token_count); - dispatch.kernel.integer_parameters.emplace("output_token_count", gather_add.output_token_count); - dispatch.kernel.integer_parameters.emplace("hidden_size", gather_add.hidden_size); - dispatch.bindings.push_back({ gather_add.first_source->id, 0, gather_add.first_source->byte_count }); - dispatch.bindings.push_back({ gather_add.second_source->id, 0, gather_add.second_source->byte_count }); - dispatch.bindings.push_back({ gather_add.row_ids->id, 0, gather_add.row_ids->byte_count }); - dispatch.bindings.push_back({ gather_add.output->id, 0, gather_add.output->byte_count }); - - match.covered_nodes.push_back(gather_add.first_get_rows_index); - match.covered_nodes.push_back(gather_add.second_get_rows_index); - match.covered_nodes.push_back(gather_add.add_node_index); - match.dispatches.push_back(std::move(dispatch)); - return true; -} - -void register_gather_add_dispatch(DispatchRegistryBuilder & registry) { - registry.add({ - "common.gather_add_f32", - GGML_OP_GET_ROWS, - DispatchMatchKind::Fused, - 1000, - DispatchSource::Common, - match_gather_add_f32_dispatch, - }); -} - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.cpp deleted file mode 100644 index eca2a207c394..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.cpp +++ /dev/null @@ -1,163 +0,0 @@ -#include "dispatch-llm-matmul.h" - -#include "dispatch-llm-shapes.h" -#include "ggml.h" -#include "graph/graph-matcher.h" -#include "kernel-corpus/kernel-corpus-catalog-verify.h" - -#include -#include -#include - -namespace ggml::hrx { -namespace { - -static constexpr KernelCatalogRef kLlmDenseLinearQ4KF16WmmaKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q4k_f16_wmma"); -static constexpr KernelCatalogRef kLlmDenseLinearQ6KF16WmmaKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q6k_f16_wmma"); - -static const Value * graph_value(const Graph & graph, ValueId id) { - return graph.values().find(id); -} - -static bool is_2d(const Value & value) { - return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; -} - -static bool is_supported_dense_input_size(int64_t input_size) { - return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; -} - -static bool is_supported_dense_output_size(int64_t output_size) { - return output_size >= 1 && output_size <= 262144; -} - -enum class LlmDenseMatmulRoute { - Q4K, - Q6K, -}; - -static std::string to_config_value(int64_t value) { - return std::to_string(value); -} - -struct LlmDenseMatmulMatch { - const Value * input = nullptr; - const Value * weight = nullptr; - const Value * output = nullptr; - KernelCatalogRef kernel = {}; - int64_t input_size = 0; - int64_t output_size = 0; - int64_t token_count = 0; - - bool matched() const { - return input != nullptr && weight != nullptr && output != nullptr && kernel.id != kUncatalogedKernelId; - } -}; - -static LlmDenseMatmulMatch match_llm_dense_matmul(const Graph & graph, - const GraphNode * node, - LlmDenseMatmulRoute route) { - LlmDenseMatmulMatch match; - if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { - return match; - } - - const Value * weight = graph_value(graph, node->inputs[0]); - const Value * input = graph_value(graph, node->inputs[1]); - const Value * output = graph_value(graph, node->output); - if (weight == nullptr || input == nullptr || output == nullptr || !is_2d(*weight) || !is_2d(*input) || - !is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || - input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { - 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 || - !is_llm_prefill_query_length(kActiveLlmMoeDispatchProfile, token_count) || - !is_supported_dense_input_size(input_size) || !is_supported_dense_output_size(output_size)) { - return {}; - } - - if (route == LlmDenseMatmulRoute::Q4K && weight->type == GGML_TYPE_Q4_K) { - match.kernel = kLlmDenseLinearQ4KF16WmmaKernel; - } else if (route == LlmDenseMatmulRoute::Q6K && weight->type == GGML_TYPE_Q6_K) { - match.kernel = kLlmDenseLinearQ6KF16WmmaKernel; - } else { - return {}; - } - - match.input = input; - match.weight = weight; - match.output = output; - match.input_size = input_size; - match.output_size = output_size; - match.token_count = token_count; - return match; -} - -} // namespace - -static void build_llm_dense_matmul_dispatch(const LlmDenseMatmulMatch & match, - DispatchMatch & dispatch_match, - size_t root_index) { - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(match.kernel); - dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", - to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", - to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "0"); - dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - 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(root_index); - dispatch_match.dispatches.push_back(std::move(dispatch)); -} - -static bool match_llm_dense_q4k_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const LlmDenseMatmulMatch match = - match_llm_dense_matmul(context.graph, context.root_node, LlmDenseMatmulRoute::Q4K); - if (!match.matched()) { - return false; - } - build_llm_dense_matmul_dispatch(match, dispatch_match, context.root_index); - return true; -} - -static bool match_llm_dense_q6k_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const LlmDenseMatmulMatch match = - match_llm_dense_matmul(context.graph, context.root_node, LlmDenseMatmulRoute::Q6K); - if (!match.matched()) { - return false; - } - build_llm_dense_matmul_dispatch(match, dispatch_match, context.root_index); - return true; -} - -void register_llm_matmul_dispatches(DispatchRegistryBuilder & registry) { - registry.add({ - "llm.matmul.dense_q4k_f16_wmma", - GGML_OP_MUL_MAT, - DispatchMatchKind::SingleOp, - 100, - DispatchSource::Llm, - match_llm_dense_q4k_dispatch, - }); - registry.add({ - "llm.matmul.dense_q6k_f16_wmma", - GGML_OP_MUL_MAT, - DispatchMatchKind::SingleOp, - 100, - DispatchSource::Llm, - match_llm_dense_q6k_dispatch, - }); -} - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.h deleted file mode 100644 index 65136aaff293..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-matmul.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "dispatch-registry.h" - -namespace ggml::hrx { - -void register_llm_matmul_dispatches(DispatchRegistryBuilder & registry); - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.cpp deleted file mode 100644 index a75509a7e88d..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.cpp +++ /dev/null @@ -1,441 +0,0 @@ -#include "dispatch-qwen-flash-attention.h" - -#include "dispatch-llm-shapes.h" -#include "ggml.h" -#include "graph/graph-matcher.h" -#include "kernel-corpus/kernel-corpus-catalog-verify.h" - -#include -#include -#include -#include - -namespace ggml::hrx { -namespace { - -static constexpr KernelCatalogRef kQwenFlashAttentionF32F16WmmaKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_flash_attention_f32_f16_wmma"); -static constexpr KernelCatalogRef kQwenFlashAttentionDecodeSplitNextQ8Kernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8"); -static constexpr int64_t kQwenAttentionHeadSize = 128; -static constexpr int64_t kQwenQueryHeadCount = 32; -static constexpr int64_t kQwenKeyValueHeadCount = 4; -static constexpr int64_t kQwenDecodeRowCapacity = 16; -static constexpr int64_t kQwenDecodeKvTileSize = 64; - -static const Value * graph_value(const Graph & graph, ValueId id) { - return graph.values().find(id); -} - -static bool nearly_equal(float lhs, float rhs) { - return std::fabs(lhs - rhs) <= 1.0e-6f; -} - -static bool is_supported_key_value_token_count(int64_t token_count) { - return token_count >= 1 && token_count <= 32768; -} - -static bool is_supported_decode_key_value_token_count(int64_t token_count) { - return token_count >= 1 && token_count <= 2048; -} - -static bool is_qwen_flash_decode_query_length(int64_t query_length) { - return query_length >= 1 && query_length < kQwenDecodeRowCapacity; -} - -static bool is_supported_head_count(int64_t head_count) { - return head_count >= 1 && head_count <= 64; -} - -static bool has_query_layout(const Value & value, int64_t query_head_count) { - const size_t element_size = sizeof(float); - return value.nb[0] == element_size && - value.nb[1] == static_cast(query_head_count * kQwenAttentionHeadSize) * element_size && - (value.ne[2] == 1 || value.nb[2] == static_cast(kQwenAttentionHeadSize) * element_size); -} - -static bool has_key_value_layout(const Value & value, int64_t key_value_head_count) { - const size_t element_size = sizeof(ggml_fp16_t); - return value.nb[0] == element_size && - value.nb[1] == static_cast(key_value_head_count * kQwenAttentionHeadSize) * element_size && - (value.ne[2] == 1 || value.nb[2] == static_cast(kQwenAttentionHeadSize) * element_size); -} - -static bool has_mask_layout(const Value & value, int64_t key_value_token_count) { - const size_t element_size = sizeof(ggml_fp16_t); - return value.nb[0] == element_size && value.nb[1] == static_cast(key_value_token_count) * element_size; -} - -static bool has_output_layout(const Value & value, int64_t query_head_count) { - const size_t element_size = sizeof(float); - return value.nb[0] == element_size && value.nb[1] == static_cast(kQwenAttentionHeadSize) * element_size && - value.nb[2] == static_cast(query_head_count * kQwenAttentionHeadSize) * element_size; -} - -static std::string to_config_value(int64_t value) { - return std::to_string(value); -} - -static size_t attention_mask_byte_count(int64_t query_token_count, int64_t key_value_token_count) { - if (query_token_count <= 0 || key_value_token_count <= 0) { - return 0; - } - return static_cast(query_token_count) * static_cast(key_value_token_count) * sizeof(ggml_fp16_t); -} - -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; - } - return static_cast(row_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); -} - -static int64_t ceil_div(int64_t value, int64_t divisor) { - return (value + divisor - 1) / divisor; -} - -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); -} - -struct QwenFlashAttentionMatch { - const Value * query = nullptr; - const Value * key = nullptr; - const Value * value = nullptr; - const Value * mask = nullptr; - const Value * output = nullptr; - const GraphNode * output_layout = nullptr; - ValueId mask_binding_value; - size_t mask_binding_bytes = 0; - int64_t query_token_count = 0; - int64_t key_value_token_count = 0; - int64_t query_head_count = 0; - int64_t key_value_head_count = 0; - - bool matched() const { - return query != nullptr && key != nullptr && value != nullptr && mask != nullptr && output != nullptr; - } -}; - -struct QwenDecodeSplitFlashAttentionMatch { - const Value * query = nullptr; - const Value * key = nullptr; - const Value * value = nullptr; - const Value * mask = nullptr; - const Value * output = nullptr; - const GraphNode * output_layout = nullptr; - int64_t query_token_count = 0; - int64_t key_value_token_count = 0; - int64_t key_value_capacity = 0; - int64_t query_head_count = 0; - int64_t key_value_head_count = 0; - - bool matched() const { - return query != nullptr && key != nullptr && value != nullptr && mask != nullptr && output != nullptr; - } -}; - -static bool has_qwen_flash_attention_params(const GraphNode & node) { - const FlashAttnExtParams * params = op_params_as(node.params); - if (params == nullptr) { - return false; - } - const float expected_scale = 1.0f / std::sqrt(static_cast(kQwenAttentionHeadSize)); - return nearly_equal(params->scale, expected_scale) && nearly_equal(params->max_bias, 0.0f) && - nearly_equal(params->logit_softcap, 0.0f) && - (params->prec == GGML_PREC_DEFAULT || params->prec == GGML_PREC_F32); -} - -static QwenFlashAttentionMatch match_qwen_flash_attention(const Graph & graph, - const CommandPlan & plan, - const GraphNode * node) { - QwenFlashAttentionMatch match; - if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4 || - !has_qwen_flash_attention_params(*node)) { - return match; - } - - const Value * query = graph_value(graph, node->inputs[0]); - const Value * key = graph_value(graph, node->inputs[1]); - const Value * value = graph_value(graph, node->inputs[2]); - const Value * mask = graph_value(graph, node->inputs[3]); - const Value * output = graph_value(graph, node->output); - if (query == nullptr || key == nullptr || value == nullptr || mask == nullptr || output == nullptr) { - return {}; - } - if (query->type != GGML_TYPE_F32 || key->type != GGML_TYPE_F16 || value->type != GGML_TYPE_F16 || - mask->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F32) { - return {}; - } - if (query->ne[0] != kQwenAttentionHeadSize || key->ne[0] != kQwenAttentionHeadSize || - value->ne[0] != kQwenAttentionHeadSize || output->ne[0] != kQwenAttentionHeadSize) { - return {}; - } - if (query->ne[3] != 1 || key->ne[3] != 1 || value->ne[3] != 1 || output->ne[3] != 1 || mask->ne[2] != 1 || - mask->ne[3] != 1) { - return {}; - } - - const int64_t query_token_count = query->ne[1]; - const int64_t query_head_count = query->ne[2]; - const int64_t key_value_capacity = key->ne[1]; - const int64_t key_value_head_count = key->ne[2]; - if (!is_qwen_prefill_query_length(query_token_count) || key_value_capacity < query_token_count || - !is_supported_head_count(query_head_count) || !is_supported_head_count(key_value_head_count) || - query_head_count % key_value_head_count != 0) { - return {}; - } - if (value->ne[1] != key_value_capacity || value->ne[2] != key_value_head_count || - mask->ne[0] > key_value_capacity || mask->ne[1] != query_token_count || output->ne[1] != query_head_count || - output->ne[2] != query_token_count) { - return {}; - } - - int64_t key_value_token_count = key_value_capacity; - ValueId mask_binding_value = mask->id; - size_t mask_binding_bytes = mask->byte_count; - const size_t compact_mask_bytes = attention_mask_byte_count(query_token_count, query_token_count); - const auto * compact_mask = find_alternate_value(plan, mask->id, GGML_TYPE_F16, compact_mask_bytes); - const bool mask_is_compact = mask->ne[0] == query_token_count; - const bool mask_is_capacity = mask->ne[0] == key_value_capacity; - if (compact_mask != nullptr && mask_is_capacity && mask->ne[0] > query_token_count) { - key_value_token_count = query_token_count; - mask_binding_value = compact_mask->alternate_value; - mask_binding_bytes = compact_mask->byte_count; - } else if (!mask_is_compact && !mask_is_capacity) { - return {}; - } - if (!is_supported_key_value_token_count(key_value_token_count)) { - return {}; - } - - if (!has_query_layout(*query, query_head_count) || !has_key_value_layout(*key, key_value_head_count) || - !has_key_value_layout(*value, key_value_head_count) || !has_output_layout(*output, query_head_count)) { - return {}; - } - if (mask_binding_value == mask->id && !has_mask_layout(*mask, key_value_token_count)) { - return {}; - } - - match.query = query; - match.key = key; - match.value = value; - match.mask = mask; - match.output = output; - match.output_layout = find_single_layout_alias_consumer(graph, output->id); - match.mask_binding_value = mask_binding_value; - match.mask_binding_bytes = mask_binding_bytes; - match.query_token_count = query_token_count; - match.key_value_token_count = key_value_token_count; - match.query_head_count = query_head_count; - match.key_value_head_count = key_value_head_count; - return match; -} - -static QwenDecodeSplitFlashAttentionMatch match_qwen_decode_split_flash_attention(const Graph & graph, - const GraphNode * node) { - QwenDecodeSplitFlashAttentionMatch match; - if (node == nullptr || node->op != GGML_OP_FLASH_ATTN_EXT || node->inputs.size() != 4 || - !has_qwen_flash_attention_params(*node)) { - return match; - } - - const Value * query = graph_value(graph, node->inputs[0]); - const Value * key = graph_value(graph, node->inputs[1]); - const Value * value = graph_value(graph, node->inputs[2]); - const Value * mask = graph_value(graph, node->inputs[3]); - const Value * output = graph_value(graph, node->output); - if (query == nullptr || key == nullptr || value == nullptr || mask == nullptr || output == nullptr) { - return {}; - } - if (query->type != GGML_TYPE_F32 || key->type != GGML_TYPE_F16 || value->type != GGML_TYPE_F16 || - mask->type != GGML_TYPE_F16 || output->type != GGML_TYPE_F32) { - return {}; - } - if (query->ne[0] != kQwenAttentionHeadSize || key->ne[0] != kQwenAttentionHeadSize || - value->ne[0] != kQwenAttentionHeadSize || output->ne[0] != kQwenAttentionHeadSize) { - return {}; - } - if (query->ne[3] != 1 || key->ne[3] != 1 || value->ne[3] != 1 || output->ne[3] != 1 || mask->ne[2] != 1 || - mask->ne[3] != 1) { - return {}; - } - - const int64_t query_token_count = query->ne[1]; - const int64_t query_head_count = query->ne[2]; - const int64_t key_value_capacity = key->ne[1]; - const int64_t key_value_head_count = key->ne[2]; - const int64_t key_value_token_count = mask->ne[0]; - if (!is_qwen_flash_decode_query_length(query_token_count) || - !is_supported_decode_key_value_token_count(key_value_token_count) || query_head_count != kQwenQueryHeadCount || - key_value_head_count != kQwenKeyValueHeadCount || key_value_capacity < key_value_token_count || - value->ne[1] != key_value_capacity || value->ne[2] != key_value_head_count || - mask->ne[1] != query_token_count || output->ne[1] != query_head_count || output->ne[2] != query_token_count) { - return {}; - } - if (!has_query_layout(*query, query_head_count) || !has_key_value_layout(*key, key_value_head_count) || - !has_key_value_layout(*value, key_value_head_count) || !has_mask_layout(*mask, key_value_token_count) || - !has_output_layout(*output, query_head_count)) { - return {}; - } - - match.query = query; - match.key = key; - match.value = value; - match.mask = mask; - match.output = output; - 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, kQwenDecodeKvTileSize) * kQwenDecodeKvTileSize; - match.query_head_count = query_head_count; - match.key_value_head_count = key_value_head_count; - return match; -} - -} // namespace - -static bool match_qwen_decode_split_flash_attention_next_q8_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { - const QwenDecodeSplitFlashAttentionMatch match = - match_qwen_decode_split_flash_attention(context.graph, context.root_node); - if (!match.matched()) { - return false; - } - - const int64_t key_value_block_count = ceil_div(match.key_value_capacity, kQwenDecodeKvTileSize); - const size_t partial_scalar_count = static_cast(match.key_value_head_count) * - static_cast(key_value_block_count) * - static_cast(kQwenDecodeRowCapacity); - const size_t partial_value_count = partial_scalar_count * static_cast(kQwenAttentionHeadSize); - const size_t partial_scalar_bytes = partial_scalar_count * sizeof(float); - const size_t partial_output_bytes = partial_value_count * sizeof(ggml_fp16_t); - const int64_t hidden_size = match.query_head_count * kQwenAttentionHeadSize; - const size_t q8_row_bytes = q8_1_x4_byte_count(1, hidden_size); - const size_t q8_output_bytes = q8_1_x4_byte_count(match.query_token_count, hidden_size); - if (partial_scalar_bytes == 0 || partial_output_bytes == 0 || q8_row_bytes == 0 || q8_output_bytes == 0) { - return false; - } - - const ValueId partial_max = match_value(context, dispatch_match, 0); - 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, "qwen.decode.flash_attention.partial_max", partial_scalar_bytes, 256 }); - dispatch_match.transients.push_back( - { partial_sum, "qwen.decode.flash_attention.partial_sum", partial_scalar_bytes, 256 }); - dispatch_match.transients.push_back( - { partial_output, "qwen.decode.flash_attention.partial_output", partial_output_bytes, 256 }); - dispatch_match.transients.push_back( - { q8_output, "qwen.decode.flash_attention.next_q8_output", q8_output_bytes, 256 }); - dispatch_match.completion_counter_requests.push_back({ - completion_counter, - "qwen.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, - "qwen.decode.flash_attention.next_q8_output" }, - metadata_status)) { - dispatch_match.status.append(metadata_status); - return false; - } - - const size_t query_row_bytes = static_cast(hidden_size) * sizeof(float); - const size_t mask_row_bytes = static_cast(match.key_value_token_count) * sizeof(ggml_fp16_t); - const size_t output_row_bytes = query_row_bytes; - for (int64_t row = 0; row < match.query_token_count; ++row) { - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kQwenFlashAttentionDecodeSplitNextQ8Kernel); - dispatch.kernel.integer_parameters.emplace("key_value_token_count", match.key_value_token_count); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.query_head_count", - to_config_value(match.query_head_count)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_head_count", - to_config_value(match.key_value_head_count)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_token_capacity", - to_config_value(match.key_value_capacity)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", "1"); - dispatch.bindings.push_back( - { match.query->id, static_cast(row) * match.query->nb[1], query_row_bytes }); - dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); - dispatch.bindings.push_back({ match.value->id, 0, match.value->byte_count }); - dispatch.bindings.push_back({ match.mask->id, static_cast(row) * match.mask->nb[1], mask_row_bytes }); - dispatch.bindings.push_back({ partial_max, 0, partial_scalar_bytes }); - dispatch.bindings.push_back({ partial_sum, 0, partial_scalar_bytes }); - dispatch.bindings.push_back({ partial_output, 0, partial_output_bytes }); - dispatch.bindings.push_back( - { 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_match.dispatches.push_back(std::move(dispatch)); - } - - dispatch_match.covered_nodes.push_back(context.root_index); - if (match.output_layout != nullptr) { - if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.output_layout, - dispatch_match.covered_nodes)) { - return false; - } - } - return true; -} - -static bool match_qwen_flash_attention_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const QwenFlashAttentionMatch match = match_qwen_flash_attention(context.graph, context.plan, context.root_node); - if (!match.matched()) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kQwenFlashAttentionF32F16WmmaKernel); - 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); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.query_head_count", - to_config_value(match.query_head_count)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_head_count", - to_config_value(match.key_value_head_count)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", - to_config_value(match.query_token_count)); - dispatch.bindings.push_back({ match.query->id, 0, match.query->byte_count }); - dispatch.bindings.push_back({ match.key->id, 0, match.key->byte_count }); - dispatch.bindings.push_back({ match.value->id, 0, match.value->byte_count }); - dispatch.bindings.push_back({ match.mask_binding_value, 0, match.mask_binding_bytes }); - dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); - - dispatch_match.covered_nodes.push_back(context.root_index); - if (match.output_layout != nullptr) { - if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.output_layout, - dispatch_match.covered_nodes)) { - return false; - } - } - dispatch_match.dispatches.push_back(std::move(dispatch)); - return true; -} - -void register_qwen_flash_attention_dispatches(DispatchRegistryBuilder & registry) { - registry.add({ - "qwen.flash_attention_decode_split_next_q8", - GGML_OP_FLASH_ATTN_EXT, - DispatchMatchKind::SingleOp, - 200, - DispatchSource::Qwen, - match_qwen_decode_split_flash_attention_next_q8_dispatch, - }); - registry.add({ - "qwen.flash_attention_f32_f16_wmma", - GGML_OP_FLASH_ATTN_EXT, - DispatchMatchKind::SingleOp, - 100, - DispatchSource::Qwen, - match_qwen_flash_attention_dispatch, - }); -} - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.h deleted file mode 100644 index 5288f3b656d2..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-flash-attention.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "dispatch-registry.h" - -namespace ggml::hrx { - -void register_qwen_flash_attention_dispatches(DispatchRegistryBuilder & registry); - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.h deleted file mode 100644 index c2f9a6bb9783..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "dispatch-registry.h" - -namespace ggml::hrx { - -void register_qwen_matmul_dispatches(DispatchRegistryBuilder & registry); - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.cpp deleted file mode 100644 index 5f5be8d6d0c5..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.cpp +++ /dev/null @@ -1,124 +0,0 @@ -#include "dispatch-qwen-preamble.h" - -#include "ggml.h" -#include "kernel-corpus/kernel-corpus-catalog-verify.h" - -#include -#include - -namespace ggml::hrx { -namespace { - -static constexpr KernelCatalogRef kQwenTokenEmbeddingQ4KKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen_token_embedding_q4k"); - -static const Value * graph_value(const Graph & graph, ValueId id) { - return graph.values().find(id); -} - -static bool is_1d_or_2d_column(const Value & value) { - return value.ne[0] > 0 && value.ne[1] == 1 && value.ne[2] == 1 && value.ne[3] == 1; -} - -static bool is_2d(const Value & value) { - return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; -} - -static bool is_supported_hidden_size(int64_t hidden_size) { - return hidden_size == 2048; -} - -static bool is_supported_token_count(int64_t token_count) { - return token_count >= 1 && token_count <= 2048; -} - -static bool is_supported_vocabulary_count(int64_t vocabulary_count) { - return vocabulary_count >= 1 && vocabulary_count <= 262144; -} - -struct QwenTokenEmbeddingMatch { - const Value * token_ids = nullptr; - const Value * weight = nullptr; - const Value * output = nullptr; - int64_t token_count = 0; - int64_t vocabulary_count = 0; - int64_t hidden_size = 0; - - bool matched() const { return token_ids != nullptr && weight != nullptr && output != nullptr; } -}; - -static QwenTokenEmbeddingMatch match_qwen_token_embedding(const Graph & graph, const GraphNode * node) { - QwenTokenEmbeddingMatch match; - if (node == nullptr || node->op != GGML_OP_GET_ROWS || node->inputs.size() != 2) { - return match; - } - - const Value * weight = graph_value(graph, node->inputs[0]); - const Value * token_ids = graph_value(graph, node->inputs[1]); - const Value * output = graph_value(graph, node->output); - if (weight == nullptr || token_ids == nullptr || output == nullptr) { - return {}; - } - if (weight->type != GGML_TYPE_Q4_K || token_ids->type != GGML_TYPE_I32 || output->type != GGML_TYPE_F32) { - return {}; - } - if (!weight->contiguous || !token_ids->contiguous || !output->contiguous) { - return {}; - } - if (!is_2d(*weight) || !is_1d_or_2d_column(*token_ids) || !is_2d(*output)) { - return {}; - } - - const int64_t hidden_size = weight->ne[0]; - const int64_t vocabulary_count = weight->ne[1]; - const int64_t token_count = token_ids->ne[0]; - if (output->ne[0] != hidden_size || output->ne[1] != token_count) { - return {}; - } - if (!is_supported_hidden_size(hidden_size) || !is_supported_vocabulary_count(vocabulary_count) || - !is_supported_token_count(token_count)) { - return {}; - } - - match.token_ids = token_ids; - match.weight = weight; - match.output = output; - match.token_count = token_count; - match.vocabulary_count = vocabulary_count; - match.hidden_size = hidden_size; - return match; -} - -} // namespace - -static bool match_qwen_token_embedding_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - const QwenTokenEmbeddingMatch match = match_qwen_token_embedding(context.graph, context.root_node); - if (!match.matched()) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kQwenTokenEmbeddingQ4KKernel); - dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.integer_parameters.emplace("vocabulary_count", match.vocabulary_count); - dispatch.bindings.push_back({ match.token_ids->id, 0, match.token_ids->byte_count }); - 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; -} - -void register_qwen_preamble_dispatches(DispatchRegistryBuilder & registry) { - registry.add({ - "qwen.preamble.token_embedding_q4k", - GGML_OP_GET_ROWS, - DispatchMatchKind::SingleOp, - 100, - DispatchSource::Qwen, - match_qwen_token_embedding_dispatch, - }); -} - -} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.h deleted file mode 100644 index 1aee68b7c3d7..000000000000 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-preamble.h +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "dispatch-registry.h" - -namespace ggml::hrx { - -void register_qwen_preamble_dispatches(DispatchRegistryBuilder & registry); - -} // 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 81bc307338e1..193f3dac977e 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp @@ -1,17 +1,14 @@ #include "dispatch-registry.h" -#include "dispatch-add.h" -#include "dispatch-gather-add.h" -#include "dispatch-llm-matmul.h" -#include "dispatch-moe-router.h" -#include "dispatch-qwen-attention-postprocess.h" -#include "dispatch-qwen-flash-attention.h" -#include "dispatch-qwen-matmul.h" -#include "dispatch-qwen-preamble.h" -#include "dispatch-rmsnorm.h" -#include "dispatch-routed-ffn.h" +#include "common/dispatch-common.h" +#include "llm/dispatch-attention-qkv.h" +#include "llm/dispatch-gated-delta-net.h" +#include "llm/dispatch-ssm-conv.h" +#include "qwen/dispatch-qwen.h" #include +#include +#include #include namespace ggml::hrx { @@ -20,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; @@ -31,22 +29,24 @@ static void sort_registrations(std::vector & registrations [](const DispatchRegistration & lhs, const DispatchRegistration & rhs) { return lhs.priority > rhs.priority; }); } -static void register_llm_dispatches(DispatchRegistryBuilder & builder) { - register_qwen_attention_postprocess_dispatches(builder); - register_qwen_flash_attention_dispatches(builder); - register_qwen_matmul_dispatches(builder); - register_llm_matmul_dispatches(builder); - register_routed_ffn_dispatches(builder); - register_qwen_preamble_dispatches(builder); - register_qwen_rmsnorm_dispatches(builder); - register_moe_router_dispatches(builder); +static bool qwen_dispatch_disabled_from_environment() { + const char * value = std::getenv("GGML_HRX_DISABLE_QWEN_DISPATCH"); + if (value == nullptr || value[0] == '\0') { + return false; + } + return std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0 && std::strcmp(value, "FALSE") != 0 && + std::strcmp(value, "off") != 0 && std::strcmp(value, "OFF") != 0; } -static DispatchRegistry build_llm_registry() { +static DispatchRegistry build_registry(bool include_qwen) { DispatchRegistryBuilder builder; - register_add_dispatch(builder); - register_gather_add_dispatch(builder); - register_llm_dispatches(builder); + register_common_dispatches(builder); + register_llm_attention_qkv_dispatches(builder); + register_llm_gated_delta_net_dispatch(builder); + register_llm_ssm_conv_dispatch(builder); + if (include_qwen) { + register_qwen_dispatches(builder); + } return builder.build(); } @@ -110,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; @@ -127,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); @@ -142,14 +165,18 @@ DispatchRegistry DispatchRegistryBuilder::build() { } const DispatchRegistry * find_dispatch_registry(const DispatchTarget & target) { - static const DispatchRegistry gfx1100_registry = build_llm_registry(); - static const DispatchRegistry gfx1151_registry = build_llm_registry(); + static const DispatchRegistry gfx1100_registry = build_registry(true); + static const DispatchRegistry gfx1100_generic_only_registry = build_registry(false); + static const DispatchRegistry gfx1151_registry = build_registry(true); + static const DispatchRegistry gfx1151_generic_only_registry = build_registry(false); + + const bool qwen_disabled = qwen_dispatch_disabled_from_environment(); if (target.architecture == "gfx1100") { - return &gfx1100_registry; + return qwen_disabled ? &gfx1100_generic_only_registry : &gfx1100_registry; } if (target.architecture == "gfx1151") { - return &gfx1151_registry; + return qwen_disabled ? &gfx1151_generic_only_registry : &gfx1151_registry; } return nullptr; } 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 new file mode 100644 index 000000000000..75256bc4ec7d --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp @@ -0,0 +1,769 @@ +#include "dispatch-attention-qkv.h" + +#include "../common/dispatch-mul-mat-common.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kAttentionQMatMulRopeTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_tiled_f32_f32"); +static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32"); +static constexpr KernelCatalogRef kAttentionVMatMulSetRowsTiledF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_tiled_f32_f32"); +static constexpr KernelCatalogRef kAttentionQMatMulRopeVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_vector_f32_f32"); +static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_vector_f32_f32"); +static constexpr KernelCatalogRef kAttentionVMatMulSetRowsVectorF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_vector_f32_f32"); +static constexpr KernelCatalogRef kAttentionQMatMulRopeLegacyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_q_matmul_rope_legacy_f32_f32"); +static constexpr KernelCatalogRef kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_k_matmul_rope_set_rows_legacy_f32_f32"); +static constexpr KernelCatalogRef kAttentionVMatMulSetRowsLegacyF32F32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_attention_v_matmul_set_rows_legacy_f32_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_1d_shape(const Value & value, int64_t ne0) { + return value.ne[0] == ne0 && value.ne[1] == 1 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_2d_shape(const Value & value, int64_t ne0, int64_t ne1) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_2d(const Value & value) { + return value.ne[0] > 0 && value.ne[1] > 0 && value.ne[2] == 1 && value.ne[3] == 1; +} + +static bool is_attention_rope_shape(const Value & value) { + return value.ne[0] >= 4 && value.ne[0] <= 1024 && value.ne[0] % 4 == 0 && value.ne[1] >= 1 && value.ne[1] <= 64 && + value.ne[2] >= 1 && value.ne[2] <= 2048 && value.ne[3] == 1; +} + +static bool is_supported_dense_input_size(int64_t input_size) { + return input_size >= 256 && input_size <= 32768 && input_size % 256 == 0; +} + +static bool is_supported_decode_input_size(int64_t input_size) { + return input_size >= 256 && input_size <= 32768 && input_size % 32 == 0; +} + +static bool is_supported_dense_output_size(int64_t output_size) { + return output_size >= 1 && output_size <= 262144; +} + +static bool is_supported_decode_output_size(int64_t output_size) { + return output_size >= 1 && output_size <= 32768; +} + +static bool is_supported_prefill_token_count(int64_t token_count) { + return token_count > 1 && token_count <= 2048; +} + +static bool is_supported_decode_token_count(int64_t token_count) { + return token_count == 1; +} + +static bool is_supported_attention_token_count(int64_t token_count) { + return is_supported_decode_token_count(token_count) || is_supported_prefill_token_count(token_count); +} + +static bool is_supported_attention_input_size(int64_t input_size, int64_t token_count) { + return is_supported_decode_token_count(token_count) ? is_supported_decode_input_size(input_size) : + is_supported_dense_input_size(input_size); +} + +static bool is_supported_attention_output_size(int64_t output_size, int64_t token_count) { + return is_supported_decode_token_count(token_count) ? is_supported_decode_output_size(output_size) : + is_supported_dense_output_size(output_size); +} + +static bool is_supported_attention_cache_row_count(int64_t row_count) { + return row_count >= 1 && row_count <= 1048576; +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (lhs.ne[i] != rhs.ne[i]) { + return false; + } + } + return true; +} + +static std::string to_config_value(int64_t value) { + return std::to_string(value); +} + +static bool is_supported_attention_rope_params(const RopeParams & params, int64_t head_size) { + return params.n_dims == head_size && params.mode == GGML_ROPE_TYPE_NORMAL && std::isfinite(params.freq_base) && + params.freq_base > 0.0f && std::isfinite(params.freq_scale) && params.freq_scale > 0.0f && + params.ext_factor == 0.0f && params.attn_factor == 1.0f; +} + +static bool build_rope_theta_table(const GraphNode & rope, int64_t head_size, std::vector & data) { + const RopeParams * params = op_params_as(rope.params); + if (params == nullptr || !is_supported_attention_rope_params(*params, head_size)) { + return false; + } + + data.resize(static_cast(head_size / 2) * sizeof(float)); + const float theta_scale = std::pow(params->freq_base, -2.0f / static_cast(head_size)); + float theta = params->freq_scale; + for (int64_t i = 0; i < head_size / 2; ++i) { + std::memcpy(data.data() + static_cast(i) * sizeof(float), &theta, sizeof(theta)); + theta *= theta_scale; + } + return true; +} + +static void build_unit_frequency_factors(int64_t head_size, std::vector & data) { + data.resize(static_cast(head_size / 2) * sizeof(float)); + const float one = 1.0f; + for (int64_t i = 0; i < head_size / 2; ++i) { + std::memcpy(data.data() + static_cast(i) * sizeof(float), &one, sizeof(one)); + } +} + +static bool cache_output_format_value(ggml_type type, int64_t & value) { + switch (type) { + case GGML_TYPE_F16: + value = 16; + return true; + case GGML_TYPE_F32: + value = 32; + return true; + default: + return false; + } +} + +static ValueId next_match_transient_value(const DispatchMatchContext & context, const DispatchMatch & dispatch_match) { + return ValueId(context.next_plan_value.value + static_cast(dispatch_match.transients.size()) + + static_cast(dispatch_match.completion_counter_requests.size())); +} + +static bool node_is_available(const DispatchMatchContext & context, const GraphNode * node) { + size_t node_index = 0; + return node != nullptr && context.graph.index().node_index(node, node_index) && + node_index < context.covered_nodes.size() && !context.covered_nodes[node_index]; +} + +static const GraphNode * find_only_available_consumer(const DispatchMatchContext & context, ValueId value) { + if (!context.graph.has_index()) { + return nullptr; + } + const std::vector & consumers = context.graph.index().consumers(value); + if (consumers.size() != 1 || !node_is_available(context, consumers.front())) { + return nullptr; + } + return consumers.front(); +} + +static const GraphNode * find_only_consumer_after_layout_aliases(const DispatchMatchContext & context, + const Value * start, + std::vector & layouts, + const Value *& current, + size_t max_layouts) { + current = start; + for (;;) { + if (current == nullptr) { + return nullptr; + } + const GraphNode * consumer = find_only_available_consumer(context, current->id); + if (consumer == nullptr) { + return nullptr; + } + if (!is_layout_alias_node(context.graph, *consumer)) { + return consumer; + } + if (layouts.size() >= max_layouts) { + return nullptr; + } + const Value * output = graph_value(context.graph, consumer->output); + if (output == nullptr || output->type != current->type || !output->contiguous) { + return nullptr; + } + layouts.push_back(consumer); + current = output; + } +} + +struct AttentionMatMulMatch { + 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; } +}; + +struct AttentionRopeMatch { + const GraphNode * node = nullptr; + const Value * input = nullptr; + const Value * positions = nullptr; + const Value * freq_factors = nullptr; + const Value * output = nullptr; + int64_t token_count = 0; + int64_t head_count = 0; + int64_t head_size = 0; + size_t theta_bytes = 0; + size_t freq_factors_bytes = 0; + std::vector theta_data; + std::vector freq_factors_data; + + bool matched() const { return node != nullptr && input != nullptr && positions != nullptr && output != nullptr; } +}; + +struct AttentionSetRowsMatch { + const GraphNode * node = nullptr; + const Value * rows = nullptr; + const Value * indices = nullptr; + const Value * output = nullptr; + int64_t output_format = 0; + int64_t token_count = 0; + int64_t cache_row_count = 0; + int64_t output_size = 0; + + bool matched() const { return node != nullptr && rows != nullptr && indices != nullptr && output != nullptr; } +}; + +enum class AttentionQkvProjectionKind { + Query, + Key, + Value, +}; + +struct AttentionQkvMatch { + AttentionMatMulMatch root; + AttentionRopeMatch rope; + AttentionSetRowsMatch set_rows; + std::vector layout_nodes; + KernelCatalogRef kernel = {}; + AttentionQkvProjectionKind kind = AttentionQkvProjectionKind::Query; + bool use_f16_activation = false; + + bool matched() const { return root.matched() && kernel.id != kUncatalogedKernelId; } +}; + +static bool prefill_v_cache_fusion_disabled() { + const char * value = std::getenv("GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"); + if (value == nullptr || value[0] == '\0') { + return false; + } + return std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0 && std::strcmp(value, "FALSE") != 0 && + std::strcmp(value, "off") != 0 && std::strcmp(value, "OFF") != 0; +} + +static bool distinct_storage(const Graph & graph, const Value & lhs, const Value & rhs) { + return !graph.values().same_storage(lhs.id, rhs.id); +} + +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 || + root.weight->alias_source.value >= 0 || !distinct_storage(context.graph, *root.input, *root.weight) || + !distinct_storage(context.graph, *root.input, *set_rows.indices) || + !distinct_storage(context.graph, *root.input, *set_rows.output) || + !distinct_storage(context.graph, *root.weight, *set_rows.indices) || + !distinct_storage(context.graph, *root.weight, *set_rows.output) || + !distinct_storage(context.graph, *set_rows.indices, *set_rows.output)) { + 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; +} + +static bool can_use_packed_lowtoken_attention_matmul(const AttentionMatMulMatch & root) { + if (root.token_count < 1 || root.token_count > 5 || root.output_size % 64 != 0 || + root.weight->alias_source.value >= 0) { + return false; + } + if (root.weight_format != CommonMulMatWeightFormat::Q4K && + root.weight_format != CommonMulMatWeightFormat::Q6K) { + return false; + } + return common_is_supported_dense_input_size(root.weight_format, root.input_size); +} + +static CommonMulMatWeightFormat packed_attention_lowtoken_format(CommonMulMatWeightFormat format) { + return format == CommonMulMatWeightFormat::Q4K ? CommonMulMatWeightFormat::Q4KRow64 : + CommonMulMatWeightFormat::Q6KRow64; +} + +static AttentionMatMulMatch match_attention_matmul_any_format(const Graph & graph, const GraphNode * node) { + AttentionMatMulMatch match; + if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { + return match; + } + + const Value * weight = graph_value(graph, node->inputs[0]); + const Value * input = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (weight == nullptr || input == nullptr || output == nullptr || !is_2d(*weight) || !is_2d(*input) || + !is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || + input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { + return {}; + } + + CommonMulMatWeightFormat format = CommonMulMatWeightFormat::Q4K; + 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 || + !is_supported_attention_token_count(token_count) || + !is_supported_attention_input_size(input_size, token_count) || + !is_supported_attention_output_size(output_size, token_count)) { + return {}; + } + + match.input = input; + match.weight = weight; + match.output = output; + match.input_size = input_size; + match.output_size = output_size; + match.token_count = token_count; + match.weight_format = format; + return match; +} + +static AttentionRopeMatch match_attention_rope(const Graph & graph, + const GraphNode * node, + const Value * expected_input) { + AttentionRopeMatch match; + if (node == nullptr || node->op != GGML_OP_ROPE || node->inputs.size() < 2 || node->inputs.size() > 3) { + return match; + } + + const Value * input = graph_value(graph, node->inputs[0]); + const Value * positions = graph_value(graph, node->inputs[1]); + const Value * output = graph_value(graph, node->output); + if (input == nullptr || positions == nullptr || output == nullptr || expected_input == nullptr || + input->id != expected_input->id || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || + positions->type != GGML_TYPE_I32 || !input->contiguous || !output->contiguous || !positions->contiguous || + !is_attention_rope_shape(*input) || !same_shape(*input, *output)) { + return {}; + } + + const int64_t head_size = input->ne[0]; + const int64_t head_count = input->ne[1]; + const int64_t token_count = input->ne[2]; + const RopeParams * params = op_params_as(node->params); + if (params == nullptr || !is_supported_attention_rope_params(*params, head_size) || + !is_1d_shape(*positions, token_count)) { + return {}; + } + + std::vector theta_data; + if (!build_rope_theta_table(*node, head_size, theta_data)) { + return {}; + } + + const Value * freq_factors = nullptr; + size_t freq_factors_bytes = 0; + std::vector freq_factors_data; + if (node->inputs.size() == 3) { + freq_factors = graph_value(graph, node->inputs[2]); + if (freq_factors == nullptr || freq_factors->type != GGML_TYPE_F32 || !freq_factors->contiguous || + !is_1d_shape(*freq_factors, head_size / 2)) { + return {}; + } + freq_factors_bytes = freq_factors->byte_count; + } else { + build_unit_frequency_factors(head_size, freq_factors_data); + freq_factors_bytes = freq_factors_data.size(); + } + + match.node = node; + match.input = input; + match.positions = positions; + match.freq_factors = freq_factors; + match.output = output; + match.token_count = token_count; + match.head_count = head_count; + match.head_size = head_size; + match.theta_bytes = theta_data.size(); + match.freq_factors_bytes = freq_factors_bytes; + match.theta_data = std::move(theta_data); + match.freq_factors_data = std::move(freq_factors_data); + return match; +} + +static AttentionSetRowsMatch match_attention_set_rows(const Graph & graph, + const GraphNode * node, + const Value * expected_rows) { + AttentionSetRowsMatch match; + if (node == nullptr || node->op != GGML_OP_SET_ROWS || node->inputs.size() != 3) { + return match; + } + + const Value * rows = graph_value(graph, node->inputs[0]); + const Value * indices = graph_value(graph, node->inputs[1]); + const Value * cache = graph_value(graph, node->inputs[2]); + const Value * output = graph_value(graph, node->output); + if (rows == nullptr || indices == nullptr || cache == nullptr || output == nullptr || expected_rows == nullptr || + rows->id != expected_rows->id || rows->type != GGML_TYPE_F32 || indices->type != GGML_TYPE_I64 || + output->type != cache->type || !rows->contiguous || !indices->contiguous || !cache->contiguous || + !same_shape(*cache, *output) || !graph.values().same_storage(cache->id, output->id)) { + return {}; + } + + int64_t output_format = 0; + if (!cache_output_format_value(output->type, output_format)) { + return {}; + } + + const int64_t output_size = rows->ne[0]; + const int64_t token_count = rows->ne[1]; + const int64_t cache_row_count = cache->ne[1]; + if (!is_2d_shape(*rows, output_size, token_count) || !is_1d_shape(*indices, token_count) || + !is_2d_shape(*cache, output_size, cache_row_count) || + !is_supported_attention_output_size(output_size, token_count) || + !is_supported_attention_token_count(token_count) || !is_supported_attention_cache_row_count(cache_row_count)) { + return {}; + } + + match.node = node; + match.rows = rows; + match.indices = indices; + match.output = output; + match.output_format = output_format; + match.token_count = token_count; + match.cache_row_count = cache_row_count; + match.output_size = output_size; + 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); + if (!root.matched() || !context.graph.has_index()) { + return match; + } + + const Value * after_projection = nullptr; + const GraphNode * consumer = + find_only_consumer_after_layout_aliases(context, root.output, match.layout_nodes, after_projection, 2); + if (consumer == nullptr) { + return {}; + } + + if (consumer->op == GGML_OP_ROPE) { + match.rope = match_attention_rope(context.graph, consumer, after_projection); + if (!match.rope.matched() || match.rope.token_count != root.token_count || + match.rope.head_size * match.rope.head_count != root.output_size) { + return {}; + } + + const Value * after_rope = nullptr; + std::vector set_rows_layouts; + const GraphNode * rope_consumer = + find_only_consumer_after_layout_aliases(context, match.rope.output, set_rows_layouts, after_rope, 2); + if (rope_consumer != nullptr && rope_consumer->op == GGML_OP_SET_ROWS) { + match.set_rows = match_attention_set_rows(context.graph, rope_consumer, after_rope); + if (!match.set_rows.matched() || match.set_rows.token_count != root.token_count || + match.set_rows.output_size != root.output_size) { + return {}; + } + match.layout_nodes.insert(match.layout_nodes.end(), set_rows_layouts.begin(), set_rows_layouts.end()); + match.root = root; + match.kernel = is_supported_decode_token_count(root.token_count) ? + kAttentionKMatMulRopeSetRowsVectorF32F32Kernel : + kAttentionKMatMulRopeSetRowsTiledF32F32Kernel; + if (can_use_packed_lowtoken_attention_matmul(root)) { + match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); + match.kernel = kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel; + } + match.kind = AttentionQkvProjectionKind::Key; + return match; + } + + match.root = root; + match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionQMatMulRopeVectorF32F32Kernel : + kAttentionQMatMulRopeTiledF32F32Kernel; + if (can_use_packed_lowtoken_attention_matmul(root)) { + match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); + match.kernel = kAttentionQMatMulRopeLegacyF32F32Kernel; + } + match.kind = AttentionQkvProjectionKind::Query; + return match; + } + + if (consumer->op == GGML_OP_SET_ROWS) { + match.set_rows = match_attention_set_rows(context.graph, consumer, after_projection); + if (!match.set_rows.matched() || match.set_rows.token_count != root.token_count || + match.set_rows.output_size != root.output_size) { + return {}; + } + const bool quantized_prefill = !is_supported_decode_token_count(root.token_count) && + !common_mul_mat_dense_float_format(root.weight_format); + if (quantized_prefill && !can_fuse_q6_f16_prefill_v_cache_cell(context, root, match.set_rows)) { + return {}; + } + match.root = root; + match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionVMatMulSetRowsVectorF32F32Kernel : + kAttentionVMatMulSetRowsTiledF32F32Kernel; + if (can_use_packed_lowtoken_attention_matmul(root)) { + match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); + match.kernel = kAttentionVMatMulSetRowsLegacyF32F32Kernel; + } + match.kind = AttentionQkvProjectionKind::Value; + match.use_f16_activation = quantized_prefill; + return match; + } + + return {}; +} + +static ValueId add_attention_constant_binding(const DispatchMatchContext & context, + DispatchMatch & dispatch_match, + const char * name, + size_t byte_count, + const std::vector & data) { + const ValueId value = next_match_transient_value(context, dispatch_match); + dispatch_match.transients.push_back({ value, name, byte_count, 256 }); + dispatch_match.constant_initializations.push_back({ + value, + name, + 0, + data, + }); + return value; +} + +static std::pair add_attention_rope_frequency_bindings(const DispatchMatchContext & context, + const AttentionRopeMatch & match, + DispatchMatch & dispatch_match) { + const ValueId theta = add_attention_constant_binding(context, dispatch_match, "llm.attention_qkv.theta", + match.theta_bytes, match.theta_data); + if (match.freq_factors != nullptr) { + return { theta, match.freq_factors->id }; + } + + const ValueId freq_factors = add_attention_constant_binding( + context, dispatch_match, "llm.attention_qkv.freq_factors", match.freq_factors_bytes, match.freq_factors_data); + return { theta, freq_factors }; +} + +static void add_attention_qkv_compile_parameters(Dispatch & dispatch, const AttentionQkvMatch & match) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", to_config_value(match.root.token_count)); + dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.input_size", to_config_value(match.root.input_size)); + dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.output_size", + to_config_value(match.root.output_size)); + dispatch.kernel.compile_parameters.emplace( + "llm.attention_qkv.weight_format", + to_config_value(common_mul_mat_format_config_value(match.root.weight_format))); + if (match.kind == AttentionQkvProjectionKind::Query || match.kind == AttentionQkvProjectionKind::Key) { + dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.head_size", + to_config_value(match.rope.head_size)); + dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.head_count", + to_config_value(match.rope.head_count)); + } + if (match.kind == AttentionQkvProjectionKind::Key || match.kind == AttentionQkvProjectionKind::Value) { + dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.cache_row_count", + to_config_value(match.set_rows.cache_row_count)); + dispatch.kernel.compile_parameters.emplace("llm.attention_qkv.cache_output_format", + to_config_value(match.set_rows.output_format)); + } +} + +static bool cover_attention_qkv_nodes(const DispatchMatchContext & context, + const AttentionQkvMatch & match, + DispatchMatch & dispatch_match) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, + dispatch_match.covered_nodes)) { + return false; + } + for (const GraphNode * layout : match.layout_nodes) { + if (!append_covered_node_index_once(context.graph, context.covered_nodes, layout, + dispatch_match.covered_nodes)) { + return false; + } + } + if ((match.kind == AttentionQkvProjectionKind::Query || match.kind == AttentionQkvProjectionKind::Key) && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.rope.node, + dispatch_match.covered_nodes)) { + return false; + } + if ((match.kind == AttentionQkvProjectionKind::Key || match.kind == AttentionQkvProjectionKind::Value) && + !append_covered_node_index_once(context.graph, context.covered_nodes, match.set_rows.node, + dispatch_match.covered_nodes)) { + return false; + } + return true; +} + +static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + AttentionQkvMatch match = match_attention_qkv_projection(context); + if (!match.matched() || !cover_attention_qkv_nodes(context, match, dispatch_match)) { + return false; + } + + Dispatch dispatch; + bool pack_q4 = match.root.weight_format == CommonMulMatWeightFormat::Q4KRow64; + bool pack_q6 = match.root.weight_format == CommonMulMatWeightFormat::Q6KRow64; + bool use_q8_activation = false; + DispatchBinding activation = { match.root.input->id, 0, match.root.input->byte_count }; + if (match.use_f16_activation) { + const size_t activation_bytes = + static_cast(match.root.input_size * match.root.token_count) * sizeof(ggml_fp16_t); + const CommandPlanAlternateValue * alternate = + find_alternate_value(context.graph, context.plan, match.root.input->id, GGML_TYPE_F16, activation_bytes); + if (alternate == nullptr) { + return false; + } + activation = { alternate->alternate_value, 0, activation_bytes }; + } + if (pack_q4 || pack_q6) { + if (common_prepare_q8_1_x4_input(context, *match.root.input, match.root.input_size, match.root.token_count, + dispatch_match, activation, + CommonQ8ActivationPolicy::ExistingAlternateOnly)) { + use_q8_activation = true; + } else { + CommonMulMatWeightFormat fallback_format = CommonMulMatWeightFormat::Q4K; + if (!common_mul_mat_format_for_type(match.root.weight->type, fallback_format)) { + return false; + } + match.root.weight_format = fallback_format; + pack_q4 = false; + pack_q6 = false; + switch (match.kind) { + case AttentionQkvProjectionKind::Query: + match.kernel = is_supported_decode_token_count(match.root.token_count) ? + kAttentionQMatMulRopeVectorF32F32Kernel : + kAttentionQMatMulRopeTiledF32F32Kernel; + break; + case AttentionQkvProjectionKind::Key: + match.kernel = is_supported_decode_token_count(match.root.token_count) ? + kAttentionKMatMulRopeSetRowsVectorF32F32Kernel : + kAttentionKMatMulRopeSetRowsTiledF32F32Kernel; + break; + case AttentionQkvProjectionKind::Value: + match.kernel = is_supported_decode_token_count(match.root.token_count) ? + kAttentionVMatMulSetRowsVectorF32F32Kernel : + kAttentionVMatMulSetRowsTiledF32F32Kernel; + break; + } + } + } + dispatch.kernel = make_kernel_specialization(match.kernel); + dispatch.kernel.integer_parameters.emplace("token_count", match.root.token_count); + add_attention_qkv_compile_parameters(dispatch, match); + if (use_q8_activation) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); + } else if (match.use_f16_activation) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); + } + dispatch.bindings.push_back(activation); + if (pack_q4 || pack_q6) { + const char * layout = pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + dispatch.bindings.push_back({ match.root.weight->id, 0, match.root.weight->byte_count, layout, + match.root.weight->type, match.root.input_size, match.root.output_size, + match.root.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.root.weight->id, 0, match.root.weight->byte_count }); + } + + if (match.kind == AttentionQkvProjectionKind::Query) { + const auto [theta, freq_factors] = add_attention_rope_frequency_bindings(context, match.rope, dispatch_match); + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + dispatch.bindings.push_back({ match.rope.output->id, 0, match.rope.output->byte_count }); + } else if (match.kind == AttentionQkvProjectionKind::Key) { + const auto [theta, freq_factors] = add_attention_rope_frequency_bindings(context, match.rope, dispatch_match); + dispatch.bindings.push_back({ match.rope.positions->id, 0, match.rope.positions->byte_count }); + dispatch.bindings.push_back({ match.set_rows.indices->id, 0, match.set_rows.indices->byte_count }); + dispatch.bindings.push_back({ theta, 0, match.rope.theta_bytes }); + dispatch.bindings.push_back({ freq_factors, 0, match.rope.freq_factors_bytes }); + dispatch.bindings.push_back({ match.set_rows.output->id, 0, match.set_rows.output->byte_count }); + } else { + dispatch.bindings.push_back({ match.set_rows.indices->id, 0, match.set_rows.indices->byte_count }); + dispatch.bindings.push_back({ match.set_rows.output->id, 0, match.set_rows.output->byte_count }); + } + + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +} // 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, + DispatchMatchKind::Fused, + 210, + DispatchSource::Llm, + match_attention_qkv_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.h b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.h new file mode 100644 index 000000000000..27884ce1468a --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.h @@ -0,0 +1,9 @@ +#pragma once + +#include "dispatch_registration/dispatch-registry.h" + +namespace ggml::hrx { + +void register_llm_attention_qkv_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..56b3888dc82f --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.cpp @@ -0,0 +1,1233 @@ +#include "dispatch-gated-delta-net.h" + +#include "../common/dispatch-mul-mat-common.h" +#include "../common/dispatch-rmsnorm.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kGatedDeltaNetProjectionEpilogueKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_projection_epilogue_f32"); +static constexpr KernelCatalogRef kGatedDeltaNetPrefillKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128"); +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_publish"); +static constexpr KernelCatalogRef kGatedDeltaNetPrefillProjectionEpilogueKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128_projection_epilogue"); +static constexpr KernelCatalogRef kGatedDeltaNetInplaceKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128_inplace"); +static constexpr KernelCatalogRef kGatedDeltaNetInplaceProjectionEpilogueKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128_inplace_projection_epilogue"); +static constexpr KernelCatalogRef kGatedDeltaNetSnapshotKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_gated_delta_net_f32_wmma_head128_snapshot"); +static constexpr KernelCatalogRef kGatedDeltaNetSnapshotProjectionEpilogueKernel = GGML_HRX_KERNEL_REF( + "loom_libs", "llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue"); +static constexpr KernelCatalogRef kGatedDeltaNetSelectedSnapshotKernel = GGML_HRX_KERNEL_REF( + "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 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"); +static constexpr KernelCatalogRef kMulMatDualQ4F32DecodeKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_dual_q4_f32_decode"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_shape(const Value & value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == ne2 && value.ne[3] == ne3; +} + +static bool is_f32(const Value * value) { + return value != nullptr && value->type == GGML_TYPE_F32; +} + +static bool distinct_storage(const Value & lhs, const Value & rhs) { + return lhs.storage != rhs.storage; +} + +static const GraphNode * producer_with_op(const Graph & graph, ValueId value, ggml_op op) { + const GraphNode * producer = graph.index().producer(value); + return producer != nullptr && producer->op == op ? producer : nullptr; +} + +static bool has_unary_op(const GraphNode * node, UnaryKind op) { + const UnaryParams * params = node != nullptr ? op_params_as(node->params) : nullptr; + return node != nullptr && node->op == GGML_OP_UNARY && params != nullptr && params->op == op; +} + +static bool append_covered_node(const DispatchMatchContext & context, const GraphNode * node, DispatchMatch & match) { + return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes); +} + +struct GatedDeltaNetMatch { + std::vector covered; + const Value * alpha_raw = nullptr; + const Value * beta_raw = nullptr; + const Value * bias = nullptr; + const Value * a_scale = nullptr; + const Value * gate = nullptr; + const Value * gate_flat = nullptr; + const Value * beta = nullptr; + const Value * raw_q = nullptr; + const Value * raw_k = nullptr; + const Value * v = nullptr; + const Value * state = nullptr; + const Value * gdn_output = nullptr; + const Value * attention = nullptr; + const Value * new_state = nullptr; + const Value * cache = nullptr; + int64_t width = 0; + int64_t q_head_count = 0; + int64_t head_count = 0; + int64_t token_count = 0; + int64_t sequence_count = 0; + int64_t snapshot_count = 0; + float l2_epsilon = 0.0f; + + bool matched() const { + return !covered.empty() && gate != nullptr && beta != nullptr && raw_q != nullptr && raw_k != nullptr && + v != nullptr && state != nullptr && gdn_output != nullptr && new_state != nullptr && cache != nullptr; + } + + bool has_projection_epilogue() const { + return alpha_raw != nullptr && beta_raw != nullptr && bias != nullptr && a_scale != nullptr && + gate_flat != nullptr; + } +}; + +struct GatedDeltaNetProjectionPairMatch { + const GraphNode * first_node = nullptr; + const GraphNode * second_node = nullptr; + const Value * input = nullptr; + const Value * first_weight = nullptr; + const Value * second_weight = nullptr; + const Value * first_output = nullptr; + const Value * second_output = nullptr; + ValueId activation; + int64_t input_size = 0; + int64_t output_size = 0; + int64_t token_count = 0; + bool native_f32 = false; + + bool matched() const { + return first_node != nullptr && second_node != nullptr && input != nullptr && first_weight != nullptr && + second_weight != nullptr && first_output != nullptr && second_output != nullptr; + } +}; + +static GatedDeltaNetMatch match_gated_delta_net(const Graph & graph, const GraphNode * node) { + GatedDeltaNetMatch match; + if (node == nullptr || node->op != GGML_OP_L2_NORM || node->inputs.size() != 1 || !graph.has_index()) { + return match; + } + + const Value * raw_q = graph_value(graph, node->inputs[0]); + const Value * q_norm = graph_value(graph, node->output); + if (!is_f32(raw_q) || !is_f32(q_norm)) { + return {}; + } + const std::vector & q_consumers = graph.index().consumers(q_norm->id); + if (q_consumers.size() != 1 || q_consumers.front() == nullptr || + q_consumers.front()->op != GGML_OP_GATED_DELTA_NET || q_consumers.front()->inputs.size() != 6 || + q_consumers.front()->inputs[0] != q_norm->id) { + return {}; + } + const GraphNode * gdn = q_consumers.front(); + + const Value * k_norm = graph_value(graph, gdn->inputs[1]); + const Value * v = graph_value(graph, gdn->inputs[2]); + const Value * gate = graph_value(graph, gdn->inputs[3]); + const Value * beta = graph_value(graph, gdn->inputs[4]); + const Value * state = graph_value(graph, gdn->inputs[5]); + const Value * gdn_output = graph_value(graph, gdn->output); + const GraphNode * k_norm_node = k_norm != nullptr ? producer_with_op(graph, k_norm->id, GGML_OP_L2_NORM) : nullptr; + const Value * raw_k = k_norm_node != nullptr && k_norm_node->inputs.size() == 1 ? + graph_value(graph, k_norm_node->inputs[0]) : + nullptr; + if (!is_f32(k_norm) || !is_f32(raw_k) || !is_f32(v) || !is_f32(gate) || !is_f32(beta) || !is_f32(state) || + !is_f32(gdn_output)) { + return {}; + } + const RmsNormParams * l2_params = op_params_as(node->params); + if (l2_params == nullptr || !std::isfinite(l2_params->eps) || l2_params->eps <= 0.0f || + !op_params_equivalent(GGML_OP_L2_NORM, node->params, k_norm_node->params)) { + return {}; + } + + const int64_t width = raw_q->ne[0]; + const int64_t q_head_count = raw_q->ne[1]; + const int64_t token_count = raw_q->ne[2]; + const int64_t sequence_count = raw_q->ne[3]; + const int64_t head_count = v->ne[1]; + if (width != 128 || q_head_count <= 0 || q_head_count > 4096 || head_count <= 0 || head_count > 4096 || + token_count < 1 || token_count > 512 || sequence_count < 1 || sequence_count > 3) { + return {}; + } + const int64_t hidden_size = width * (2 * q_head_count + head_count); + if (!is_shape(*raw_k, width, q_head_count, token_count, sequence_count) || + !is_shape(*q_norm, width, q_head_count, token_count, sequence_count) || + !is_shape(*k_norm, width, q_head_count, token_count, sequence_count) || + !is_shape(*v, width, head_count, token_count, sequence_count) || + !is_shape(*gate, 1, head_count, token_count, sequence_count) || + !is_shape(*beta, 1, head_count, token_count, sequence_count) || + !is_shape(*state, width, width, head_count, sequence_count) || gdn_output->ne[0] != width * head_count || + gdn_output->ne[2] != 1 || gdn_output->ne[3] != 1) { + return {}; + } + const int64_t attention_rows = token_count * sequence_count; + const int64_t snapshot_rows = width * sequence_count; + if (gdn_output->ne[1] <= attention_rows || (gdn_output->ne[1] - attention_rows) % snapshot_rows != 0) { + return {}; + } + const int64_t snapshot_count = (gdn_output->ne[1] - attention_rows) / snapshot_rows; + if (snapshot_count < 1 || snapshot_count > 5) { + return {}; + } + if (!q_norm->contiguous || !k_norm->contiguous || !state->contiguous || !gdn_output->contiguous || + raw_q->storage != raw_k->storage || raw_q->storage != v->storage || raw_q->storage_offset != 0 || + raw_k->storage_offset != static_cast(width * q_head_count) * sizeof(float) || + v->storage_offset != static_cast(2 * width * q_head_count) * sizeof(float) || + raw_q->nb[0] != sizeof(float) || raw_q->nb[1] != static_cast(width) * sizeof(float) || + raw_q->nb[2] != static_cast(hidden_size) * sizeof(float) || raw_k->nb != raw_q->nb || + v->nb[0] != sizeof(float) || v->nb[1] != static_cast(width) * sizeof(float) || + v->nb[2] != static_cast(hidden_size) * sizeof(float)) { + return {}; + } + + const GraphNode * gate_reshape = producer_with_op(graph, gate->id, GGML_OP_RESHAPE); + const Value * gate_flat = gate_reshape != nullptr && gate_reshape->inputs.size() == 1 ? + graph_value(graph, gate_reshape->inputs[0]) : + nullptr; + const GraphNode * gate_mul = gate_flat != nullptr ? producer_with_op(graph, gate_flat->id, GGML_OP_MUL) : nullptr; + if (gate_mul == nullptr || gate_mul->inputs.size() != 2) { + return {}; + } + const Value * alpha_softplus = graph_value(graph, gate_mul->inputs[0]); + const Value * a_scale = graph_value(graph, gate_mul->inputs[1]); + const GraphNode * softplus = + alpha_softplus != nullptr ? producer_with_op(graph, alpha_softplus->id, GGML_OP_UNARY) : nullptr; + if (!has_unary_op(softplus, UnaryKind::SoftPlus) || softplus->inputs.size() != 1) { + return {}; + } + const Value * alpha_biased = graph_value(graph, softplus->inputs[0]); + const GraphNode * add_bias = + alpha_biased != nullptr ? producer_with_op(graph, alpha_biased->id, GGML_OP_ADD) : nullptr; + if (add_bias == nullptr || add_bias->inputs.size() != 2) { + return {}; + } + const Value * alpha = graph_value(graph, add_bias->inputs[0]); + const Value * bias = graph_value(graph, add_bias->inputs[1]); + const GraphNode * alpha_reshape = alpha != nullptr ? producer_with_op(graph, alpha->id, GGML_OP_RESHAPE) : nullptr; + const Value * alpha_raw = alpha_reshape != nullptr && alpha_reshape->inputs.size() == 1 ? + graph_value(graph, alpha_reshape->inputs[0]) : + nullptr; + + const GraphNode * sigmoid = producer_with_op(graph, beta->id, GGML_OP_UNARY); + if (!has_unary_op(sigmoid, UnaryKind::Sigmoid) || sigmoid->inputs.size() != 1) { + return {}; + } + const Value * beta_pre = graph_value(graph, sigmoid->inputs[0]); + const GraphNode * beta_reshape = + beta_pre != nullptr ? producer_with_op(graph, beta_pre->id, GGML_OP_RESHAPE) : nullptr; + const Value * beta_raw = beta_reshape != nullptr && beta_reshape->inputs.size() == 1 ? + graph_value(graph, beta_reshape->inputs[0]) : + nullptr; + if (!is_f32(alpha_raw) || !is_f32(beta_raw) || !is_f32(alpha) || !is_f32(alpha_biased) || !is_f32(alpha_softplus) || + !is_f32(a_scale) || !is_f32(bias) || !is_f32(gate_flat) || !is_f32(beta_pre) || + !is_shape(*alpha_raw, head_count, token_count * sequence_count, 1, 1) || + !is_shape(*beta_raw, head_count, token_count * sequence_count, 1, 1) || + !is_shape(*gate_flat, head_count, token_count, sequence_count, 1) || !is_shape(*bias, head_count, 1, 1, 1) || + !is_shape(*a_scale, head_count, 1, 1, 1) || !alpha_raw->contiguous || !beta_raw->contiguous || + !bias->contiguous || !a_scale->contiguous || !gate_flat->contiguous || !beta->contiguous) { + return {}; + } + + const std::array epilogue_values = { alpha_raw, beta_raw, bias, a_scale, gate_flat, beta }; + for (size_t lhs = 0; lhs < epilogue_values.size(); ++lhs) { + for (size_t rhs = lhs + 1; rhs < epilogue_values.size(); ++rhs) { + if (!distinct_storage(*epilogue_values[lhs], *epilogue_values[rhs])) { + return {}; + } + } + } + + const GraphNode * state_reshape = producer_with_op(graph, state->id, GGML_OP_RESHAPE); + if (state_reshape == nullptr || state_reshape->inputs.size() != 1) { + return {}; + } + const std::vector & gdn_consumers = graph.index().consumers(gdn_output->id); + if (gdn_consumers.size() != 2) { + return {}; + } + const GraphNode * new_state_view = nullptr; + const GraphNode * attention_view = nullptr; + const size_t attention_bytes = + static_cast(width * head_count * token_count * sequence_count) * sizeof(float); + for (const GraphNode * consumer : gdn_consumers) { + const Value * output = + consumer != nullptr && consumer->op == GGML_OP_VIEW ? graph_value(graph, consumer->output) : nullptr; + if (output != nullptr && output->storage == gdn_output->storage && + output->storage_offset == gdn_output->storage_offset + attention_bytes) { + new_state_view = consumer; + } else if (output != nullptr && output->storage == gdn_output->storage && + output->storage_offset == gdn_output->storage_offset && + is_shape(*output, width, head_count, token_count, sequence_count)) { + attention_view = consumer; + } else { + return {}; + } + } + const Value * new_state = new_state_view != nullptr ? graph_value(graph, new_state_view->output) : nullptr; + if (new_state == nullptr || attention_view == nullptr) { + return {}; + } + const int64_t written_snapshot_count = std::min(token_count, snapshot_count); + const size_t state_bytes = static_cast(width * width * head_count) * sizeof(float); + const size_t state_plane_bytes = state_bytes * static_cast(sequence_count); + const bool single_snapshot_state = + snapshot_count == 1 && is_shape(*new_state, width, width, head_count, sequence_count); + const bool rollback_snapshot_state = + is_shape(*new_state, width * width * head_count, sequence_count, written_snapshot_count, 1); + if ((!single_snapshot_state && !rollback_snapshot_state) || !new_state->contiguous || + new_state->storage != gdn_output->storage || + new_state->storage_offset != gdn_output->storage_offset + attention_bytes || + new_state->byte_count != state_plane_bytes * static_cast(written_snapshot_count)) { + return {}; + } + const std::vector & state_consumers = graph.index().consumers(new_state->id); + if (state_consumers.size() != 1 || state_consumers.front() == nullptr || + state_consumers.front()->op != GGML_OP_CPY || state_consumers.front()->inputs.size() != 2 || + state_consumers.front()->inputs[0] != new_state->id) { + return {}; + } + const GraphNode * cache_copy = state_consumers.front(); + const Value * cache_target = graph_value(graph, cache_copy->inputs[1]); + const Value * cache = graph_value(graph, cache_copy->output); + const GraphNode * cache_view = cache_target != nullptr ? graph.index().producer(cache_target->id) : nullptr; + if (!is_f32(cache_target) || !is_f32(cache) || cache_view == nullptr) { + return {}; + } + const size_t snapshot_stride_count = static_cast(written_snapshot_count - 1); + if (snapshot_stride_count != 0 && + cache->nb[2] > (std::numeric_limits::max() - state_bytes) / snapshot_stride_count) { + return {}; + } + const size_t required_cache_bytes = state_plane_bytes + snapshot_stride_count * cache->nb[2]; + if (cache_view->op != GGML_OP_VIEW || cache_view->inputs.size() != 1 || + !same_full_value_range(*cache_target, *cache) || cache->ne[0] != width * width * head_count || + cache->ne[1] != sequence_count || cache->ne[2] != written_snapshot_count || cache->ne[3] != 1 || + cache->nb[0] != sizeof(float) || cache->nb[1] != state_bytes || + (written_snapshot_count > 1 && cache->nb[2] < state_plane_bytes) || cache->byte_count < required_cache_bytes || + !distinct_storage(*new_state, *cache)) { + return {}; + } + + if (!distinct_storage(*gdn_output, *state) || !distinct_storage(*gdn_output, *gate_flat) || + !distinct_storage(*gdn_output, *beta)) { + return {}; + } + + match.covered = { alpha_reshape, add_bias, softplus, gate_mul, gate_reshape, + beta_reshape, sigmoid, node, k_norm_node, state_reshape, + gdn, new_state_view, cache_view, cache_copy, attention_view }; + match.alpha_raw = alpha_raw; + match.beta_raw = beta_raw; + match.bias = bias; + match.a_scale = a_scale; + match.gate = gate; + match.gate_flat = gate_flat; + match.beta = beta; + match.raw_q = raw_q; + match.raw_k = raw_k; + match.v = v; + match.state = state; + match.gdn_output = gdn_output; + match.attention = graph_value(graph, attention_view->output); + match.new_state = new_state; + match.cache = cache; + match.width = width; + match.q_head_count = q_head_count; + match.head_count = head_count; + match.token_count = token_count; + match.sequence_count = sequence_count; + match.snapshot_count = snapshot_count; + match.l2_epsilon = l2_params->eps; + return match; +} + +static GatedDeltaNetProjectionPairMatch match_gated_delta_net_projection_pair(const DispatchMatchContext & context) { + GatedDeltaNetProjectionPairMatch match; + const GraphNode * first_node = context.root_node; + if (first_node == nullptr || first_node->op != GGML_OP_MUL_MAT || first_node->inputs.size() != 2 || + !context.graph.has_index()) { + return match; + } + + const GraphNode * reshape = common_find_only_consumer_with_op(context.graph, first_node->output, GGML_OP_RESHAPE); + const GraphNode * add = + reshape != nullptr ? common_find_only_consumer_with_op(context.graph, reshape->output, GGML_OP_ADD) : nullptr; + const GraphNode * softplus = + add != nullptr ? common_find_only_consumer_with_op(context.graph, add->output, GGML_OP_UNARY) : nullptr; + const GraphNode * mul = has_unary_op(softplus, UnaryKind::SoftPlus) ? + common_find_only_consumer_with_op(context.graph, softplus->output, GGML_OP_MUL) : + nullptr; + const GraphNode * gate_reshape = + mul != nullptr ? common_find_only_consumer_with_op(context.graph, mul->output, GGML_OP_RESHAPE) : nullptr; + const GraphNode * gdn = + gate_reshape != nullptr ? + common_find_only_consumer_with_op(context.graph, gate_reshape->output, GGML_OP_GATED_DELTA_NET) : + nullptr; + if (gdn == nullptr || gdn->inputs.size() != 6) { + return {}; + } + + const GraphNode * q_norm = producer_with_op(context.graph, gdn->inputs[0], GGML_OP_L2_NORM); + const GatedDeltaNetMatch gdn_match = match_gated_delta_net(context.graph, q_norm); + if (!gdn_match.matched() || !gdn_match.has_projection_epilogue() || gdn_match.alpha_raw->id != first_node->output) { + return {}; + } + + const GraphNode * second_node = producer_with_op(context.graph, gdn_match.beta_raw->id, GGML_OP_MUL_MAT); + if (second_node == nullptr || second_node == first_node || second_node->inputs.size() != 2 || + second_node->inputs[1] != first_node->inputs[1]) { + return {}; + } + + CommonMulMatMatch first = common_match_mul_mat_any_format( + context.graph, first_node, kMulMatSymmetricI4LowRowAdjacentDualWmmaKernel, false); + CommonMulMatMatch second = common_match_mul_mat_any_format( + context.graph, second_node, kMulMatSymmetricI4LowRowAdjacentDualWmmaKernel, false); + if (!first.matched()) { + first = common_match_mul_mat_any_format( + context.graph, first_node, kMulMatDualQ4F32DecodeKernel, true); + } + if (!second.matched()) { + second = common_match_mul_mat_any_format( + context.graph, second_node, kMulMatDualQ4F32DecodeKernel, true); + } + if (!first.matched() || !second.matched() || first.weight->type != GGML_TYPE_Q4_K || + second.weight->type != GGML_TYPE_Q4_K || first.weight->alias_source.value >= 0 || + second.weight->alias_source.value >= 0 || first.input->id != second.input->id || + first.input_size != second.input_size || first.output_size != second.output_size || + first.token_count != second.token_count || first.output_size != gdn_match.head_count || + first.token_count != gdn_match.token_count * gdn_match.sequence_count || + !distinct_storage(*first.weight, *second.weight) || !distinct_storage(*first.output, *second.output)) { + return {}; + } + + const CommonSymmetricI4ActivationLayout activation_layout = + common_symmetric_i4_activation_layout(first.input_size, first.token_count); + const CommandPlanAlternateValue * alternate = find_alternate_value(context.graph, context.plan, first.input->id, + GGML_TYPE_COUNT, activation_layout.total_bytes); + const bool symmetric_input = first.token_count > 1 && + alternate != nullptr && alternate->name == kCommonSymmetricI4K32ActivationAlternateName; + const bool native_f32 = is_f32(first.input) && first.input->contiguous && + first.token_count == 1 && first.output_size >= 24 && first.output_size < 64 && + first.input_size >= 4096 && first.input_size <= 32768 && first.input_size % 1024 == 0; + if (!symmetric_input && !native_f32) { + return {}; + } + if (!symmetric_input) { + size_t second_index = 0; + if (!context.graph.index().node_index(second_node, second_index) || second_index <= context.root_index) { + return {}; + } + for (const Value * input : { first.input, second.weight }) { + const GraphNode * producer = context.graph.index().producer(input->storage_root); + size_t index = 0; + if (producer != nullptr && (!context.graph.index().node_index(producer, index) || + index >= context.covered_nodes.size() || !context.covered_nodes[index])) { + return {}; + } + } + const auto & nodes = context.graph.nodes(); + for (size_t i = context.root_index + 1; i < second_index; ++i) { + const GraphNode & node = nodes[i]; + if (node.op == GGML_OP_NONE || node.op == GGML_OP_VIEW || node.op == GGML_OP_RESHAPE || + node.op == GGML_OP_PERMUTE || node.op == GGML_OP_TRANSPOSE) { + continue; + } + const Value * output = graph_value(context.graph, node.output); + if (output != nullptr && (output->storage == first.input->storage || + output->storage == second.weight->storage)) { + return {}; + } + } + } + + match.first_node = first_node; + match.second_node = second_node; + match.input = first.input; + match.first_weight = first.weight; + match.second_weight = second.weight; + match.first_output = first.output; + match.second_output = second.output; + match.activation = symmetric_input ? alternate->alternate_value : ValueId{}; + match.native_f32 = !symmetric_input; + match.input_size = first.input_size; + match.output_size = first.output_size; + match.token_count = first.token_count; + return match; +} + +static GatedDeltaNetMatch match_direct_gated_delta_net(const Graph & graph, const GraphNode * gdn) { + GatedDeltaNetMatch match; + if (gdn == nullptr || gdn->op != GGML_OP_GATED_DELTA_NET || gdn->inputs.size() != 6 || !graph.has_index()) { + return match; + } + + const Value * q_norm = graph_value(graph, gdn->inputs[0]); + const Value * k_norm = graph_value(graph, gdn->inputs[1]); + const Value * v = graph_value(graph, gdn->inputs[2]); + const Value * gate = graph_value(graph, gdn->inputs[3]); + const Value * beta = graph_value(graph, gdn->inputs[4]); + const Value * state = graph_value(graph, gdn->inputs[5]); + const Value * gdn_output = graph_value(graph, gdn->output); + if (!is_f32(q_norm) || !is_f32(k_norm) || !is_f32(v) || !is_f32(gate) || !is_f32(beta) || !is_f32(state) || + !is_f32(gdn_output)) { + return {}; + } + + const int64_t width = q_norm->ne[0]; + const int64_t q_head_count = q_norm->ne[1]; + const int64_t token_count = q_norm->ne[2]; + const int64_t sequence_count = q_norm->ne[3]; + const int64_t head_count = v->ne[1]; + if (width != 128 || q_head_count <= 0 || q_head_count > 4096 || head_count <= 0 || head_count > 4096 || + token_count < 1 || token_count > 512 || sequence_count < 1 || sequence_count > 3) { + return {}; + } + + if (!is_shape(*q_norm, width, q_head_count, token_count, sequence_count) || + !is_shape(*k_norm, width, q_head_count, token_count, sequence_count) || + !is_shape(*v, width, head_count, token_count, sequence_count) || + !is_shape(*gate, 1, head_count, token_count, sequence_count) || + !is_shape(*beta, 1, head_count, token_count, sequence_count) || + !is_shape(*state, width, width, head_count, sequence_count) || gdn_output->ne[0] != width * head_count || + gdn_output->ne[2] != 1 || gdn_output->ne[3] != 1 || !q_norm->contiguous || !k_norm->contiguous || + !gate->contiguous || !beta->contiguous || !state->contiguous || !gdn_output->contiguous || + q_norm->nb[0] != sizeof(float) || q_norm->nb[1] != static_cast(width) * sizeof(float) || + q_norm->nb[2] != static_cast(width * q_head_count) * sizeof(float) || k_norm->nb != q_norm->nb || + v->nb[0] != sizeof(float) || v->nb[1] != static_cast(width) * sizeof(float) || + v->nb[2] < static_cast(width * head_count) * sizeof(float) || v->nb[2] % sizeof(float) != 0 || + v->nb[3] % sizeof(float) != 0) { + return {}; + } + + const int64_t attention_rows = token_count * sequence_count; + const int64_t snapshot_rows = width * sequence_count; + if (gdn_output->ne[1] <= attention_rows || (gdn_output->ne[1] - attention_rows) % snapshot_rows != 0) { + return {}; + } + const int64_t snapshot_count = (gdn_output->ne[1] - attention_rows) / snapshot_rows; + if (snapshot_count < 1 || snapshot_count > 5) { + return {}; + } + + const std::vector & gdn_consumers = graph.index().consumers(gdn_output->id); + if (gdn_consumers.size() != 2) { + return {}; + } + const GraphNode * attention_view = nullptr; + const GraphNode * new_state_view = nullptr; + const size_t attention_bytes = + static_cast(width * head_count * token_count * sequence_count) * sizeof(float); + for (const GraphNode * consumer : gdn_consumers) { + const Value * output = + consumer != nullptr && consumer->op == GGML_OP_VIEW ? graph_value(graph, consumer->output) : nullptr; + if (output != nullptr && output->storage == gdn_output->storage && + output->storage_offset == gdn_output->storage_offset && + is_shape(*output, width, head_count, token_count, sequence_count)) { + attention_view = consumer; + } else if (output != nullptr && output->storage == gdn_output->storage && + output->storage_offset == gdn_output->storage_offset + attention_bytes) { + new_state_view = consumer; + } else { + return {}; + } + } + + const Value * new_state = new_state_view != nullptr ? graph_value(graph, new_state_view->output) : nullptr; + if (attention_view == nullptr || new_state == nullptr) { + return {}; + } + const int64_t written_snapshot_count = std::min(token_count, snapshot_count); + const size_t state_bytes = static_cast(width * width * head_count) * sizeof(float); + const size_t state_plane_bytes = state_bytes * static_cast(sequence_count); + const bool single_snapshot_state = + snapshot_count == 1 && is_shape(*new_state, width, width, head_count, sequence_count); + const bool rollback_snapshot_state = + is_shape(*new_state, width * width * head_count, sequence_count, written_snapshot_count, 1); + if ((!single_snapshot_state && !rollback_snapshot_state) || !new_state->contiguous || + new_state->storage != gdn_output->storage || + new_state->storage_offset != gdn_output->storage_offset + attention_bytes || + new_state->byte_count != state_plane_bytes * static_cast(written_snapshot_count)) { + return {}; + } + + const std::vector & state_consumers = graph.index().consumers(new_state->id); + if (state_consumers.size() != 1 || state_consumers.front() == nullptr || + state_consumers.front()->op != GGML_OP_CPY || state_consumers.front()->inputs.size() != 2 || + state_consumers.front()->inputs[0] != new_state->id) { + return {}; + } + const GraphNode * cache_copy = state_consumers.front(); + const Value * cache_target = graph_value(graph, cache_copy->inputs[1]); + const Value * cache = graph_value(graph, cache_copy->output); + const GraphNode * cache_view = cache_target != nullptr ? graph.index().producer(cache_target->id) : nullptr; + if (!is_f32(cache_target) || !is_f32(cache) || + (cache_view != nullptr && (cache_view->op != GGML_OP_VIEW || cache_view->inputs.size() != 1)) || + !same_full_value_range(*cache_target, *cache) || cache->ne[0] != width * width * head_count || + cache->ne[1] != sequence_count || cache->ne[2] != written_snapshot_count || cache->ne[3] != 1 || + cache->nb[0] != sizeof(float) || cache->nb[1] != state_bytes || + (written_snapshot_count > 1 && cache->nb[2] < state_plane_bytes) || !distinct_storage(*new_state, *cache) || + !distinct_storage(*gdn_output, *state) || !distinct_storage(*gdn_output, *gate) || + !distinct_storage(*gdn_output, *beta)) { + return {}; + } + const size_t snapshot_stride_count = static_cast(written_snapshot_count - 1); + if (snapshot_stride_count != 0 && + cache->nb[2] > (std::numeric_limits::max() - state_bytes) / snapshot_stride_count) { + return {}; + } + const size_t required_cache_bytes = state_plane_bytes + snapshot_stride_count * cache->nb[2]; + if (cache->byte_count < required_cache_bytes) { + return {}; + } + + match.covered = { gdn, attention_view, new_state_view, cache_copy }; + match.gate = gate; + match.beta = beta; + match.raw_q = q_norm; + match.raw_k = k_norm; + match.v = v; + match.state = state; + match.gdn_output = gdn_output; + match.attention = graph_value(graph, attention_view->output); + match.new_state = new_state; + match.cache = cache; + match.width = width; + match.q_head_count = q_head_count; + match.head_count = head_count; + match.token_count = token_count; + match.sequence_count = sequence_count; + match.snapshot_count = snapshot_count; + match.l2_epsilon = 1.0e-6f; + return match; +} + +static void set_compile_parameter(KernelSpecialization & kernel, const char * name, int64_t value) { + kernel.compile_parameters.emplace(name, std::to_string(value)); +} + +static void set_float_compile_parameter(KernelSpecialization & kernel, const char * name, float value) { + std::ostringstream stream; + stream << std::setprecision(std::numeric_limits::max_digits10) << value; + kernel.compile_parameters.emplace(name, stream.str()); +} + +static void configure_gated_delta_net_kernel(Dispatch & dispatch, + const GatedDeltaNetMatch & match, + int64_t token_count) { + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.head_width", match.width); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.head_count", match.head_count); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.token_count", token_count); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.sequence_count", match.sequence_count); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.qk_stride1", match.raw_q->nb[1] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.qk_stride2", match.raw_q->nb[2] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.qk_stride3", match.raw_q->nb[3] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.value_stride1", match.v->nb[1] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.value_stride2", match.v->nb[2] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.value_stride3", match.v->nb[3] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.scalar_stride1", match.beta->nb[1] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.scalar_stride2", match.beta->nb[2] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.scalar_stride3", match.beta->nb[3] / sizeof(float)); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.query_head_count", match.q_head_count); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.query_sequence_ratio", 1); + set_float_compile_parameter(dispatch.kernel, "llm.gated_delta_net.l2_epsilon", match.l2_epsilon); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.workgroup_size", 256); +} + +static bool match_gated_delta_net_projection_pair_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { + const GatedDeltaNetProjectionPairMatch match = match_gated_delta_net_projection_pair(context); + if (!match.matched() || !append_covered_node(context, match.first_node, dispatch_match) || + !append_covered_node(context, match.second_node, dispatch_match)) { + return false; + } + + if (match.native_f32) { + Dispatch projections; + projections.kernel = make_kernel_specialization(kMulMatDualQ4F32DecodeKernel); + set_compile_parameter(projections.kernel, "ggml.mul_mat_dual_q4_f32_c1.input_size", match.input_size); + set_compile_parameter(projections.kernel, "ggml.mul_mat_dual_q4_f32_c1.output_size", match.output_size); + projections.bindings = { + { match.input->id, 0, match.input->byte_count }, + { match.first_weight->id, 0, match.first_weight->byte_count }, + { match.second_weight->id, 0, match.second_weight->byte_count }, + { match.first_output->id, 0, match.first_output->byte_count }, + { match.second_output->id, 0, match.second_output->byte_count }, + }; + dispatch_match.dispatches.push_back(std::move(projections)); + return true; + } + + const CommonSymmetricI4ActivationLayout activation_layout = + common_symmetric_i4_activation_layout(match.input_size, match.token_count); + + Dispatch projections; + projections.kernel = make_kernel_specialization(kMulMatSymmetricI4LowRowAdjacentDualWmmaKernel); + projections.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.input_size", + common_to_config_value(match.input_size)); + projections.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.output_size", + common_to_config_value(match.output_size)); + projections.kernel.compile_parameters.emplace("ggml.mul_mat.symmetric_i4.lowrow.token_count", + common_to_config_value(match.token_count)); + projections.kernel.compile_parameters.emplace( + "ggml.mul_mat.symmetric_i4.lowrow.row_group_size", + common_to_config_value(static_cast( + common_symmetric_shared4_row_group_size(match.input_size, match.output_size, 4)))); + projections.bindings.push_back( + common_symmetric_i4_shared4_weight_binding(*match.first_weight, match.input_size, match.output_size)); + projections.bindings.push_back( + common_symmetric_i4_shared4_weight_binding(*match.second_weight, match.input_size, match.output_size)); + projections.bindings.push_back({ match.first_output->id, 0, match.first_output->byte_count }); + projections.bindings.push_back({ match.second_output->id, 0, match.second_output->byte_count }); + projections.bindings.push_back({ match.activation, 0, activation_layout.payload_bytes }); + projections.bindings.push_back( + { match.activation, activation_layout.scales_offset, activation_layout.metadata_bytes }); + projections.bindings.push_back( + { match.activation, activation_layout.sums_offset, activation_layout.metadata_bytes }); + dispatch_match.dispatches.push_back(std::move(projections)); + return true; +} + +static bool match_gated_delta_net_rmsnorm_gate(const DispatchMatchContext & context, + const GatedDeltaNetMatch & gdn, + DispatchMatch & match) { + if (gdn.attention->kind != ValueKind::Transient || gdn.gdn_output->kind != ValueKind::Transient) { + return false; + } + const GraphNode * rms = common_find_only_consumer_with_op(context.graph, gdn.attention->id, GGML_OP_RMS_NORM); + DispatchMatchContext rms_context = context; + rms_context.root_node = rms; + if (rms == nullptr || !context.graph.index().node_index(rms, rms_context.root_index) || + !common_match_rmsnorm_gate_dispatch(rms_context, match) || match.dispatches.size() != 1) { + return false; + } + 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; + } + // These inputs remain live while the fused operation publishes its rows. + const Value * output = graph_value(context.graph, norm.bindings[3].value); + for (const Value * input : {gdn.raw_q, gdn.raw_k, gdn.v, gdn.state, gdn.gate, gdn.beta, gdn.cache}) { + if (!distinct_storage(*output, *input)) { + return false; + } + } + return true; +} + +static bool match_gated_delta_net_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const GatedDeltaNetMatch match = context.root_node != nullptr && context.root_node->op == GGML_OP_GATED_DELTA_NET ? + match_direct_gated_delta_net(context.graph, context.root_node) : + match_gated_delta_net(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + const bool can_fuse_projection_epilogue = match.has_projection_epilogue() && match.token_count <= 16; + DispatchMatch rms_match; + const bool fuse_rmsnorm_gate = match.snapshot_count == 1 && !can_fuse_projection_epilogue && + match_gated_delta_net_rmsnorm_gate(context, match, rms_match); + Dispatch rms_gate; + if (fuse_rmsnorm_gate) { + rms_gate = std::move(rms_match.dispatches.front()); + rms_match.dispatches.clear(); + dispatch_match = std::move(rms_match); + } + + for (const GraphNode * node : match.covered) { + if (!append_covered_node(context, node, dispatch_match)) { + return false; + } + } + + const int64_t written_snapshot_count = std::min(match.token_count, match.snapshot_count); + const int64_t prefix_token_count = match.token_count - written_snapshot_count; + const size_t cache_stride = match.cache->nb[2]; + const bool fuse_snapshot_projection_epilogue = + can_fuse_projection_epilogue && match.snapshot_count != 1 && prefix_token_count == 0 && + written_snapshot_count == match.token_count && (match.token_count >= 2 || match.sequence_count > 1) && + match.token_count <= 5 && cache_stride % sizeof(float) == 0; + + 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_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 }); + epilogue.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + epilogue.bindings.push_back({ match.a_scale->id, 0, match.a_scale->byte_count }); + epilogue.bindings.push_back({ match.gate_flat->id, 0, match.gate_flat->byte_count }); + epilogue.bindings.push_back({ match.beta->id, 0, match.beta->byte_count }); + dispatch_match.dispatches.push_back(std::move(epilogue)); + } + + if (match.snapshot_count == 1) { + Dispatch gdn; + gdn.kernel = make_kernel_specialization(fuse_rmsnorm_gate ? kGatedDeltaNetRmsNormGateKernel : + can_fuse_projection_epilogue ? + kGatedDeltaNetPrefillProjectionEpilogueKernel : + kGatedDeltaNetPrefillKernel); + configure_gated_delta_net_kernel(gdn, match, match.token_count); + gdn.bindings.push_back({ match.raw_q->id, 0, match.raw_q->byte_count }); + gdn.bindings.push_back({ match.raw_k->id, 0, match.raw_k->byte_count }); + gdn.bindings.push_back({ match.v->id, 0, match.v->byte_count }); + if (can_fuse_projection_epilogue) { + gdn.bindings.push_back({ match.alpha_raw->id, 0, match.alpha_raw->byte_count }); + gdn.bindings.push_back({ match.beta_raw->id, 0, match.beta_raw->byte_count }); + gdn.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + gdn.bindings.push_back({ match.a_scale->id, 0, match.a_scale->byte_count }); + } else { + gdn.bindings.push_back({ match.gate->id, 0, match.gate->byte_count }); + gdn.bindings.push_back({ match.beta->id, 0, match.beta->byte_count }); + } + gdn.bindings.push_back({ match.state->id, 0, match.state->byte_count }); + gdn.bindings.push_back({ match.gdn_output->id, 0, match.gdn_output->byte_count }); + if (fuse_rmsnorm_gate) { + for (const char * parameter : {"ggml.rmsnorm_gate_f32.rms_epsilon", "ggml.rmsnorm_gate_f32.gate_op"}) { + gdn.kernel.compile_parameters.emplace(parameter, rms_gate.kernel.compile_parameters.at(parameter)); + } + gdn.bindings.insert(gdn.bindings.end(), rms_gate.bindings.begin() + 1, rms_gate.bindings.end()); + } + + Dispatch cache_copy; + cache_copy.kernel = make_kernel_specialization(kCopyF32Kernel); + cache_copy.kernel.integer_parameters.emplace("element_count", match.new_state->element_count); + cache_copy.bindings.push_back({ match.new_state->id, 0, match.new_state->byte_count }); + cache_copy.bindings.push_back({ match.cache->id, 0, match.cache->byte_count }); + + dispatch_match.dispatches.push_back(std::move(gdn)); + dispatch_match.dispatches.push_back(std::move(cache_copy)); + return true; + } + + const size_t q_token_bytes = static_cast(match.width * match.q_head_count) * sizeof(float); + const size_t v_token_bytes = static_cast(match.width * match.head_count) * sizeof(float); + const size_t gate_token_bytes = static_cast(match.head_count) * sizeof(float); + const size_t attention_token_bytes = v_token_bytes; + const size_t state_bytes = static_cast(match.width * match.width * match.head_count) * sizeof(float); + const size_t state_plane_bytes = state_bytes * static_cast(match.sequence_count); + + if (prefix_token_count == 0 && written_snapshot_count == match.token_count && + (match.token_count >= 2 || match.sequence_count > 1) && match.token_count <= 5 && + cache_stride % sizeof(float) == 0) { + Dispatch gdn; + gdn.kernel = make_kernel_specialization(fuse_snapshot_projection_epilogue ? + kGatedDeltaNetSnapshotProjectionEpilogueKernel : + kGatedDeltaNetSnapshotKernel); + configure_gated_delta_net_kernel(gdn, match, match.token_count); + set_compile_parameter(gdn.kernel, "llm.gated_delta_net.snapshot_stride", cache_stride / sizeof(float)); + gdn.bindings.push_back({ match.raw_q->id, 0, match.raw_q->byte_count }); + gdn.bindings.push_back({ match.raw_k->id, 0, match.raw_k->byte_count }); + gdn.bindings.push_back({ match.v->id, 0, match.v->byte_count }); + if (fuse_snapshot_projection_epilogue) { + gdn.bindings.push_back({ match.alpha_raw->id, 0, match.alpha_raw->byte_count }); + gdn.bindings.push_back({ match.beta_raw->id, 0, match.beta_raw->byte_count }); + gdn.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + gdn.bindings.push_back({ match.a_scale->id, 0, match.a_scale->byte_count }); + } else { + gdn.bindings.push_back({ match.gate->id, 0, match.gate->byte_count }); + gdn.bindings.push_back({ match.beta->id, 0, match.beta->byte_count }); + } + gdn.bindings.push_back({ match.state->id, 0, state_plane_bytes }); + gdn.bindings.push_back( + { match.cache->id, 0, state_plane_bytes + static_cast(written_snapshot_count - 1) * cache_stride }); + gdn.bindings.push_back( + { match.gdn_output->id, 0, + static_cast(match.token_count * match.sequence_count) * attention_token_bytes }); + dispatch_match.dispatches.push_back(std::move(gdn)); + return true; + } + + auto append_state_copy = [&](ValueId source, size_t source_offset, ValueId target, size_t target_offset) { + Dispatch copy; + copy.kernel = make_kernel_specialization(kCopyF32Kernel); + copy.kernel.integer_parameters.emplace("element_count", state_bytes / sizeof(float)); + copy.bindings.push_back({ source, source_offset, state_bytes }); + copy.bindings.push_back({ target, target_offset, state_bytes }); + dispatch_match.dispatches.push_back(std::move(copy)); + }; + + if (prefix_token_count > 0) { + const size_t q_span = static_cast(prefix_token_count - 1) * match.raw_q->nb[2] + q_token_bytes; + const size_t k_span = static_cast(prefix_token_count - 1) * match.raw_k->nb[2] + q_token_bytes; + const size_t v_span = static_cast(prefix_token_count - 1) * match.v->nb[2] + v_token_bytes; + const size_t gate_span = static_cast(prefix_token_count - 1) * match.gate->nb[2] + gate_token_bytes; + const size_t beta_span = static_cast(prefix_token_count - 1) * match.beta->nb[2] + gate_token_bytes; + const size_t prefix_attention_bytes = static_cast(prefix_token_count) * attention_token_bytes; + + Dispatch prefix; + prefix.kernel = make_kernel_specialization(can_fuse_projection_epilogue ? + kGatedDeltaNetPrefillProjectionEpilogueKernel : + kGatedDeltaNetPrefillKernel); + configure_gated_delta_net_kernel(prefix, match, prefix_token_count); + prefix.bindings.push_back({ match.raw_q->id, 0, q_span }); + prefix.bindings.push_back({ match.raw_k->id, 0, k_span }); + prefix.bindings.push_back({ match.v->id, 0, v_span }); + if (can_fuse_projection_epilogue) { + const size_t alpha_span = static_cast(prefix_token_count - 1) * match.alpha_raw->nb[1] + + gate_token_bytes; + const size_t beta_raw_span = static_cast(prefix_token_count - 1) * match.beta_raw->nb[1] + + gate_token_bytes; + prefix.bindings.push_back({ match.alpha_raw->id, 0, alpha_span }); + prefix.bindings.push_back({ match.beta_raw->id, 0, beta_raw_span }); + prefix.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + prefix.bindings.push_back({ match.a_scale->id, 0, match.a_scale->byte_count }); + } else { + prefix.bindings.push_back({ match.gate->id, 0, gate_span }); + prefix.bindings.push_back({ match.beta->id, 0, beta_span }); + } + prefix.bindings.push_back({ match.state->id, 0, state_bytes }); + prefix.bindings.push_back({ match.gdn_output->id, 0, prefix_attention_bytes + state_bytes }); + dispatch_match.dispatches.push_back(std::move(prefix)); + append_state_copy(match.gdn_output->id, prefix_attention_bytes, match.cache->id, + static_cast(written_snapshot_count - 1) * cache_stride); + } else { + append_state_copy(match.state->id, 0, match.cache->id, + static_cast(written_snapshot_count - 1) * cache_stride); + } + + for (int64_t snapshot = 0; snapshot < written_snapshot_count; ++snapshot) { + const int64_t token = prefix_token_count + snapshot; + const int64_t slot = written_snapshot_count - 1 - snapshot; + if (snapshot > 0) { + append_state_copy(match.cache->id, static_cast(slot + 1) * cache_stride, match.cache->id, + static_cast(slot) * cache_stride); + } + + Dispatch gdn; + gdn.kernel = make_kernel_specialization(can_fuse_projection_epilogue ? + kGatedDeltaNetInplaceProjectionEpilogueKernel : + kGatedDeltaNetInplaceKernel); + configure_gated_delta_net_kernel(gdn, match, 1); + gdn.bindings.push_back({ match.raw_q->id, static_cast(token) * match.raw_q->nb[2], q_token_bytes }); + gdn.bindings.push_back({ match.raw_k->id, static_cast(token) * match.raw_k->nb[2], q_token_bytes }); + gdn.bindings.push_back({ match.v->id, static_cast(token) * match.v->nb[2], v_token_bytes }); + if (can_fuse_projection_epilogue) { + gdn.bindings.push_back( + { match.alpha_raw->id, static_cast(token) * match.alpha_raw->nb[1], gate_token_bytes }); + gdn.bindings.push_back( + { match.beta_raw->id, static_cast(token) * match.beta_raw->nb[1], gate_token_bytes }); + gdn.bindings.push_back({ match.bias->id, 0, match.bias->byte_count }); + gdn.bindings.push_back({ match.a_scale->id, 0, match.a_scale->byte_count }); + } else { + gdn.bindings.push_back( + { match.gate->id, static_cast(token) * match.gate->nb[2], gate_token_bytes }); + gdn.bindings.push_back( + { match.beta->id, static_cast(token) * match.beta->nb[2], gate_token_bytes }); + } + gdn.bindings.push_back({ match.cache->id, static_cast(slot) * cache_stride, state_bytes }); + gdn.bindings.push_back( + { match.gdn_output->id, static_cast(token) * attention_token_bytes, attention_token_bytes }); + dispatch_match.dispatches.push_back(std::move(gdn)); + } + + return true; +} +static bool match_gated_delta_net_rmsnorm_q8(const DispatchMatchContext & context, + const GatedDeltaNetMatch & gdn, + const Value & state_cache, + DispatchMatch & result) { + if (gdn.head_count < 24 || gdn.attention->kind != ValueKind::Transient || !gdn.attention->contiguous) { + return false; + } + const Graph & graph = context.graph; + const Value * input = gdn.attention; + const GraphNode * reshape = nullptr; + const GraphNode * rms = common_find_only_consumer_with_op(graph, input->id, GGML_OP_RMS_NORM); + if (rms == nullptr) { + reshape = common_find_only_consumer_with_op(graph, input->id, GGML_OP_RESHAPE); + const Value * reshaped = reshape != nullptr ? graph_value(graph, reshape->output) : nullptr; + if (reshaped == nullptr || !is_layout_alias_node(graph, *reshape) || !reshaped->contiguous || + !same_full_value_range(*input, *reshaped)) { + return false; + } + input = reshaped; + rms = common_find_only_consumer_with_op(graph, input->id, GGML_OP_RMS_NORM); + } + DispatchMatchContext rms_context = context; + rms_context.root_node = rms; + DispatchMatch match; + if (rms == nullptr || !graph.index().node_index(rms, rms_context.root_index) || + !common_match_rmsnorm_gate_dispatch(rms_context, match) || match.dispatches.size() != 1) { + return false; + } + const Dispatch & norm = match.dispatches.front(); + 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; + } + const Value * output = graph_value(graph, norm.bindings[3].value); + if (output == nullptr || output->kind != ValueKind::Transient || !distinct_storage(*output, state_cache)) { + return false; + } + for (const Value * value : {gdn.raw_q, gdn.raw_k, gdn.v, gdn.alpha_raw, gdn.beta_raw, + gdn.bias, gdn.a_scale, gdn.state, gdn.cache}) { + if (!distinct_storage(*output, *value)) { + return false; + } + } + for (size_t binding : {size_t{1}, size_t{2}}) { + const Value * value = graph_value(graph, norm.bindings[binding].value); + if (value == nullptr || !distinct_storage(*value, *gdn.cache) || !distinct_storage(*value, *output)) { + return false; + } + const GraphNode * producer = graph.index().producer(value->storage_root); + size_t index = 0; + if (producer != nullptr && (!graph.index().node_index(producer, index) || + index >= context.covered_nodes.size() || !context.covered_nodes[index])) { + return false; + } + } + if (reshape != nullptr && !append_covered_node(context, reshape, match)) { + return false; + } + result = std::move(match); + return true; +} + +static bool match_selected_gated_delta_net_dispatch(const DispatchMatchContext & context, DispatchMatch & result) { + const Graph & graph = context.graph; + const GraphNode * gather = context.root_node; + if (gather == nullptr || gather->op != GGML_OP_GET_ROWS || gather->inputs.size() != 2 || !graph.has_index()) { + return false; + } + const Value * cache = graph_value(graph, gather->inputs[0]); + const Value * ids = graph_value(graph, gather->inputs[1]); + const Value * gathered = graph_value(graph, gather->output); + if (!is_f32(cache) || !is_f32(gathered) || ids == nullptr || ids->type != GGML_TYPE_I32 || + !cache->contiguous || !ids->contiguous || !gathered->contiguous || + cache->ne[0] <= 0 || cache->ne[1] < 1 || cache->ne[1] > 262208 || + cache->ne[2] != 1 || cache->ne[3] != 1 || cache->element_count > (int64_t{1} << 30) || + !is_shape(*ids, 1, 1, 1, 1) || !is_shape(*gathered, cache->ne[0], 1, 1, 1) || + gathered->kind != ValueKind::Transient || !graph.index().has_single_consumer(gathered->id)) { + return false; + } + const GraphNode * reshape = graph.index().consumers(gathered->id).front(); + if (reshape == nullptr || reshape->op != GGML_OP_RESHAPE || reshape->inputs.size() != 1 || + !graph.index().has_single_consumer(reshape->output)) { + return false; + } + const GraphNode * gdn = graph.index().consumers(reshape->output).front(); + if (gdn == nullptr || gdn->op != GGML_OP_GATED_DELTA_NET || gdn->inputs.size() != 6 || + gdn->inputs[5] != reshape->output) { + return false; + } + const GraphNode * q_norm = producer_with_op(graph, gdn->inputs[0], GGML_OP_L2_NORM); + const GatedDeltaNetMatch match = match_gated_delta_net(graph, q_norm); + if (!match.matched() || !match.has_projection_epilogue() || match.sequence_count != 1 || + match.token_count < 1 || match.token_count > 5 || match.snapshot_count < match.token_count || + (match.token_count == 1 && match.snapshot_count == 1) || + cache->ne[0] != match.width * match.width * match.head_count || + !same_full_value_range(*gathered, *match.state) || match.gdn_output->kind != ValueKind::Transient) { + return false; + } + const size_t state_bytes = static_cast(cache->ne[0]) * sizeof(float); + if (match.cache->nb[2] % state_bytes != 0) { + return false; + } + if (cache->storage == match.cache->storage && + (match.cache->storage_offset < cache->storage_offset || + (match.cache->storage_offset - cache->storage_offset) % state_bytes != 0 || + match.cache->storage_offset - cache->storage_offset > cache->byte_count || + match.cache->byte_count > cache->byte_count - (match.cache->storage_offset - cache->storage_offset))) { + return false; + } + const auto internal = [&](const GraphNode * node) { + return std::find(match.covered.begin(), match.covered.end(), node) != match.covered.end(); + }; + const std::array inputs = { + match.raw_q, match.raw_k, match.v, match.alpha_raw, match.beta_raw, match.bias, match.a_scale, cache, ids + }; + for (const Value * input : inputs) { + const GraphNode * producer = graph.index().producer(input->storage_root); + size_t index = 0; + if (producer != nullptr && (!graph.index().node_index(producer, index) || + index >= context.covered_nodes.size() || !context.covered_nodes[index])) { + return false; + } + if (!distinct_storage(*input, *match.gdn_output) || + (input != cache && !distinct_storage(*input, *match.cache))) { + return false; + } + } + size_t copy_index = 0; + for (const GraphNode * node : match.covered) { + const Value * output = graph_value(graph, node->output); + if (node->op == GGML_OP_CPY) { + if (!graph.index().node_index(node, copy_index)) { + return false; + } + } + if (output->id == match.cache->id || output->id == match.attention->id) { + continue; + } + if (!is_layout_alias_node(graph, *node) && output->kind != ValueKind::Transient) { + return false; + } + for (const GraphNode * consumer : graph.index().consumers(output->id)) { + if (!internal(consumer)) { + return false; + } + } + } + if (copy_index <= context.root_index) { + return false; + } + // The fused dispatch publishes snapshots at the gather's position. + for (size_t index = context.root_index + 1; index < copy_index; ++index) { + const GraphNode & node = graph.nodes()[index]; + const Value * output = graph_value(graph, node.output); + if (internal(&node) || is_layout_alias_node(graph, node) || output->byte_count == 0) { + continue; + } + const auto touches_cache = [&](const Value * value) { + return value != nullptr && value->byte_count != 0 && value->storage == match.cache->storage; + }; + if (touches_cache(output)) { + return false; + } + for (ValueId input : node.inputs) { + if (touches_cache(graph_value(graph, input))) { + return false; + } + } + } + DispatchMatchContext gdn_context = context; + gdn_context.root_node = q_norm; + if (!graph.index().node_index(q_norm, gdn_context.root_index)) { + return false; + } + DispatchMatch selected; + DispatchMatch rms_match; + const bool fuse_rms_q8 = match_gated_delta_net_rmsnorm_q8(context, match, *cache, rms_match); + if ((match.token_count == 1 && !fuse_rms_q8) || + !match_gated_delta_net_dispatch(gdn_context, selected) || + selected.dispatches.size() != (match.token_count == 1 ? 2 : 1) || + selected.dispatches.back().bindings.size() != (match.token_count == 1 ? 9 : 10) || + !append_covered_node(context, gather, selected)) { + return false; + } + if (fuse_rms_q8) { + Dispatch dispatch = std::move(selected.dispatches.back()); + const Dispatch & norm = rms_match.dispatches.front(); + const auto snapshot = dispatch.bindings[match.token_count == 1 ? 7 : 8]; + auto kernel = make_kernel_specialization(kGatedDeltaNetSelectedRmsQ8Kernel); + kernel.compile_parameters = std::move(dispatch.kernel.compile_parameters); + set_compile_parameter(kernel, "llm.gated_delta_net.state_row_count", cache->ne[1]); + set_compile_parameter(kernel, "llm.gated_delta_net.snapshot_stride", match.cache->nb[2] / sizeof(float)); + for (const auto & parameter : norm.kernel.compile_parameters) { + kernel.compile_parameters.emplace(parameter); + } + dispatch.kernel = std::move(kernel); + dispatch.bindings.resize(7); + dispatch.bindings.push_back({ cache->id, 0, cache->byte_count }); + dispatch.bindings.push_back(snapshot); + dispatch.bindings.push_back(norm.bindings[3]); + dispatch.bindings.push_back({ ids->id, 0, ids->byte_count }); + dispatch.bindings.push_back(norm.bindings[1]); + dispatch.bindings.push_back(norm.bindings[2]); + dispatch.bindings.push_back(norm.bindings[4]); + selected.dispatches.clear(); + selected.dispatches.push_back(std::move(dispatch)); + selected.covered_nodes.insert(selected.covered_nodes.end(), rms_match.covered_nodes.begin(), rms_match.covered_nodes.end()); + selected.transients = std::move(rms_match.transients); + if (!selected.metadata.append(std::move(rms_match.metadata), selected.status)) { + return false; + } + result = std::move(selected); + return true; + } + Dispatch & dispatch = selected.dispatches.front(); + auto kernel = make_kernel_specialization(kGatedDeltaNetSelectedSnapshotKernel); + kernel.compile_parameters = std::move(dispatch.kernel.compile_parameters); + dispatch.kernel = std::move(kernel); + set_compile_parameter(dispatch.kernel, "llm.gated_delta_net.state_row_count", cache->ne[1]); + dispatch.bindings[7] = { cache->id, 0, cache->byte_count }; + dispatch.bindings.push_back({ ids->id, 0, ids->byte_count }); + result = std::move(selected); + return true; +} +} // namespace + +void register_llm_gated_delta_net_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "llm.gated_delta_net.selected_snapshot.f32_wmma_head128", + GGML_OP_GET_ROWS, + DispatchMatchKind::Fused, + 250, + DispatchSource::Llm, + match_selected_gated_delta_net_dispatch, + }); + registry.add({ + "llm.gated_delta_net.projection_pair.q4_k", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 300, + DispatchSource::Llm, + match_gated_delta_net_projection_pair_dispatch, + }); + registry.add({ + "llm.gated_delta_net.direct.f32_wmma_head128", + GGML_OP_GATED_DELTA_NET, + DispatchMatchKind::Fused, + 200, + DispatchSource::Llm, + match_gated_delta_net_dispatch, + }); + registry.add({ + "llm.gated_delta_net.f32_wmma_head128", + GGML_OP_L2_NORM, + DispatchMatchKind::Fused, + 200, + DispatchSource::Llm, + match_gated_delta_net_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.h b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.h new file mode 100644 index 000000000000..a02c2d2c1489 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-gated-delta-net.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_llm_gated_delta_net_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx 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 new file mode 100644 index 000000000000..a7e5a4ddb3eb --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp @@ -0,0 +1,878 @@ +#include "dispatch-ssm-conv.h" + +#include "../common/dispatch-binary-common.h" +#include "../common/dispatch-mul-mat-common.h" +#include "ggml.h" +#include "graph/graph-matcher.h" +#include "kernel-corpus/kernel-corpus-catalog-verify.h" + +#include +#include +#include +#include +#include +#include + +namespace ggml::hrx { +namespace { + +static constexpr KernelCatalogRef kMulMatConv4Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_quantized_f16_wmma_prefill_conv4"); +static constexpr KernelCatalogRef kMulMatConv4InteriorKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior"); +static constexpr KernelCatalogRef kSsmConvFinishKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_dconv4_silu_prefill_finish_f32"); +static constexpr KernelCatalogRef kSsmConvSnapshotKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_snapshot_window_tail_f32"); +static constexpr KernelCatalogRef kSsmConvPrefillKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_dconv4_silu_prefill_512_wg1024"); +static constexpr KernelCatalogRef kSsmConvDecodeKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_dconv4_silu_decode_f32"); +static constexpr KernelCatalogRef kSsmConvRollbackKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_dconv4_silu_rollback_f32"); +static constexpr KernelCatalogRef kSsmConvGenericKernel = GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_f32"); +static constexpr KernelCatalogRef kSsmConvGenericBinaryKernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_binary_f32"); +static constexpr KernelCatalogRef kLfmDconv3Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "llm_ssm_conv_lfm_dconv3_state_f32"); + +static const Value * graph_value(const Graph & graph, ValueId id) { + return graph.values().find(id); +} + +static bool is_f32(const Value * value) { + return value != nullptr && value->type == GGML_TYPE_F32; +} + +static bool is_shape(const Value & value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return value.ne[0] == ne0 && value.ne[1] == ne1 && value.ne[2] == ne2 && value.ne[3] == ne3; +} + +static bool same_shape(const Value & lhs, const Value & rhs) { + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (lhs.ne[dim] != rhs.ne[dim]) { + return false; + } + } + return true; +} + +static bool packed_f32_layout(const Value & value) { + size_t expected_stride = sizeof(float); + for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { + if (value.ne[dim] <= 0 || value.nb[dim] != expected_stride) { + return false; + } + expected_stride *= static_cast(value.ne[dim]); + } + return true; +} + +static bool supported_hidden_size(int64_t hidden_size) { + return hidden_size >= 32 && hidden_size <= 65536 && hidden_size % 32 == 0; +} + +static const Value * root_alias_value(const Graph & graph, const Value * value) { + for (size_t depth = 0; value != nullptr && value->alias_source.value >= 0 && depth < graph.values().size(); + ++depth) { + value = graph_value(graph, value->alias_source); + } + return value; +} + +static bool distinct_storage(const Value & lhs, const Value & rhs) { + return lhs.storage != rhs.storage; +} + +static bool ranges_overlap(const Value & lhs, const Value & rhs) { + if (lhs.storage != rhs.storage || lhs.byte_count == 0 || rhs.byte_count == 0) { + return false; + } + if (lhs.storage_offset <= rhs.storage_offset) { + return rhs.storage_offset - lhs.storage_offset < lhs.byte_count; + } + return lhs.storage_offset - rhs.storage_offset < rhs.byte_count; +} + +static bool append_covered_node(const DispatchMatchContext & context, const GraphNode * node, DispatchMatch & match) { + return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes); +} + +static void set_compile_parameter(KernelSpecialization & kernel, const char * name, int64_t value) { + kernel.compile_parameters.emplace(name, std::to_string(value)); +} + +struct SsmConvCacheUpdate { + const GraphNode * state_tail_view = nullptr; + const GraphNode * cache_view = nullptr; + const GraphNode * cache_copy = nullptr; + const Value * state_tail = nullptr; + const Value * cache = nullptr; + int64_t source_row = 0; +}; + +struct SsmConvPrefillMatch { + const GraphNode * concat = nullptr; + const GraphNode * ssm = nullptr; + const GraphNode * silu = nullptr; + const Value * state = nullptr; + const Value * x = nullptr; + const Value * filter = nullptr; + const Value * output = nullptr; + std::vector cache_updates; + int64_t hidden_size = 0; + int64_t token_count = 0; + int64_t sequence_count = 0; + + bool matched() const { + return concat != nullptr && ssm != nullptr && silu != nullptr && state != nullptr && x != nullptr && + filter != nullptr && output != nullptr && !cache_updates.empty(); + } +}; + +struct SsmConvCoreMatch { + const GraphNode * ssm = nullptr; + const Value * window = nullptr; + const Value * filter = nullptr; + const Value * conv_output = nullptr; + const GraphNode * unary = nullptr; + const GraphNode * binary = nullptr; + const Value * binary_operand = nullptr; + const Value * output = nullptr; + int64_t d_conv = 0; + int64_t d_inner = 0; + int64_t token_count = 0; + int64_t sequence_count = 0; + UnaryKind unary_op = UnaryKind::Identity; + BinaryKind binary_op = BinaryKind::Mul; + bool binary_lhs = true; + + bool matched() const { + return ssm != nullptr && window != nullptr && filter != nullptr && conv_output != nullptr && output != nullptr; + } + + bool has_unary_fusion() const { return unary != nullptr; } + + bool has_binary_fusion() const { return binary != nullptr; } +}; + +struct LfmDconv3Match { + const GraphNode * concat = nullptr; + const GraphNode * transpose = nullptr; + const GraphNode * ssm = nullptr; + const GraphNode * tail_view = nullptr; + const GraphNode * cache_view = nullptr; + const GraphNode * cache_copy = nullptr; + const Value * state = nullptr; + const Value * x = nullptr; + const Value * filter = nullptr; + const Value * output = nullptr; + const Value * cache = nullptr; + int64_t token_count = 0; + + bool matched() const { + return concat != nullptr && transpose != nullptr && ssm != nullptr && tail_view != nullptr && + cache_view != nullptr && cache_copy != nullptr && state != nullptr && x != nullptr && + filter != nullptr && output != nullptr && cache != nullptr; + } +}; + +static bool lfm_dconv3_disabled() { + return std::getenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION") != nullptr; +} + +static LfmDconv3Match match_lfm_dconv3(const Graph & graph, const GraphNode * node) { + LfmDconv3Match match; + if (lfm_dconv3_disabled() || node == nullptr || node->op != GGML_OP_CONCAT || node->inputs.size() != 2 || + !graph.has_index()) { + return match; + } + + const Value * state = graph_value(graph, node->inputs[0]); + const Value * x_transposed = graph_value(graph, node->inputs[1]); + const Value * window = graph_value(graph, node->output); + if (!is_f32(state) || !is_f32(x_transposed) || !is_f32(window)) { + return {}; + } + const GraphNode * transpose = graph.index().producer(x_transposed->id); + if (transpose == nullptr || transpose->op != GGML_OP_TRANSPOSE || transpose->inputs.size() != 1) { + return {}; + } + const Value * x_layout = graph_value(graph, transpose->inputs[0]); + const Value * x = root_alias_value(graph, x_layout); + if (!is_f32(x_layout) || !x_layout->contiguous || !is_f32(x) || !x->contiguous) { + return {}; + } + + const int64_t token_count = x_layout->ne[1]; + if ((token_count != 1 && token_count != 64) || !is_shape(*state, 2, 2048, 1, 1) || + !is_shape(*x_layout, 2048, token_count, 1, 1) || !is_shape(*x_transposed, token_count, 2048, 1, 1) || + !is_shape(*window, token_count + 2, 2048, 1, 1) || state->nb[0] != sizeof(float) || + state->nb[1] != 2 * sizeof(float) || x_layout->nb[0] != sizeof(float) || + x_layout->nb[1] != 2048 * sizeof(float) || x_transposed->nb[0] != 2048 * sizeof(float) || + x_transposed->nb[1] != sizeof(float) || window->nb[0] != sizeof(float) || + window->nb[1] != static_cast(token_count + 2) * sizeof(float)) { + return {}; + } + + const GraphNode * ssm = nullptr; + const GraphNode * tail_view = nullptr; + const auto & consumers = graph.index().consumers(window->id); + if (consumers.size() != 2) { + return {}; + } + for (const GraphNode * consumer : consumers) { + if (consumer != nullptr && consumer->op == GGML_OP_SSM_CONV && ssm == nullptr) { + ssm = consumer; + } else if (consumer != nullptr && consumer->op == GGML_OP_VIEW && tail_view == nullptr) { + tail_view = consumer; + } else { + return {}; + } + } + if (ssm == nullptr || tail_view == nullptr || ssm->inputs.size() != 2 || ssm->inputs[0] != window->id || + tail_view->inputs.size() != 1) { + return {}; + } + const Value * filter = graph_value(graph, ssm->inputs[1]); + const Value * output = graph_value(graph, ssm->output); + const Value * tail = graph_value(graph, tail_view->output); + if (!is_f32(filter) || !is_f32(output) || !is_f32(tail) || !filter->contiguous || !output->contiguous || + !is_shape(*filter, 3, 2048, 1, 1) || !is_shape(*output, 2048, token_count, 1, 1) || + !is_shape(*tail, 2, 2048, 1, 1) || tail->alias_source != window->id || + tail->storage_offset < window->storage_offset || + tail->storage_offset - window->storage_offset != static_cast(token_count) * sizeof(float)) { + return {}; + } + const auto & tail_consumers = graph.index().consumers(tail->id); + if (tail_consumers.size() != 1 || tail_consumers.front() == nullptr || tail_consumers.front()->op != GGML_OP_CPY || + tail_consumers.front()->inputs.size() != 2 || tail_consumers.front()->inputs[0] != tail->id) { + return {}; + } + const GraphNode * cache_copy = tail_consumers.front(); + const Value * cache_target = graph_value(graph, cache_copy->inputs[1]); + const Value * cache = graph_value(graph, cache_copy->output); + const GraphNode * cache_view = cache_target != nullptr ? graph.index().producer(cache_target->id) : nullptr; + if (!is_f32(cache_target) || !is_f32(cache) || cache_view == nullptr || cache_view->op != GGML_OP_VIEW || + cache_view->inputs.size() != 1 || !same_full_value_range(*cache_target, *cache) || + cache->byte_count != static_cast(4096) * sizeof(float)) { + return {}; + } + if (!distinct_storage(*state, *x) || !distinct_storage(*state, *filter) || !distinct_storage(*state, *output) || + !distinct_storage(*state, *cache) || !distinct_storage(*x, *filter) || !distinct_storage(*x, *output) || + !distinct_storage(*x, *cache) || !distinct_storage(*filter, *output) || !distinct_storage(*filter, *cache) || + !distinct_storage(*output, *cache)) { + return {}; + } + + match.concat = node; + match.transpose = transpose; + match.ssm = ssm; + match.tail_view = tail_view; + match.cache_view = cache_view; + match.cache_copy = cache_copy; + match.state = state; + match.x = x; + match.filter = filter; + match.output = output; + match.cache = cache; + match.token_count = token_count; + return match; +} + +static bool match_lfm_dconv3_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const LfmDconv3Match match = match_lfm_dconv3(context.graph, context.root_node); + if (!match.matched() || !append_covered_node(context, match.concat, dispatch_match) || + !append_covered_node(context, match.ssm, dispatch_match) || + !append_covered_node(context, match.tail_view, dispatch_match) || + !append_covered_node(context, match.cache_view, dispatch_match) || + !append_covered_node(context, match.cache_copy, dispatch_match)) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(kLfmDconv3Kernel); + set_compile_parameter(dispatch.kernel, "llm.ssm_conv.lfm_dconv3.n_t", match.token_count); + dispatch.bindings.push_back({match.state->id, 0, match.state->byte_count}); + dispatch.bindings.push_back({match.x->id, 0, match.x->byte_count}); + dispatch.bindings.push_back({match.filter->id, 0, match.filter->byte_count}); + dispatch.bindings.push_back({match.output->id, 0, match.output->byte_count}); + dispatch.bindings.push_back({match.cache->id, 0, match.cache->byte_count}); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static SsmConvPrefillMatch match_ssm_conv_prefill(const Graph & graph, const GraphNode * node) { + SsmConvPrefillMatch match; + if (node == nullptr || node->op != GGML_OP_CONCAT || node->inputs.size() != 2 || !graph.has_index()) { + return match; + } + + const Value * state = graph_value(graph, node->inputs[0]); + const Value * x_transposed = graph_value(graph, node->inputs[1]); + const Value * window = graph_value(graph, node->output); + if (!is_f32(state) || !is_f32(x_transposed) || !is_f32(window)) { + return {}; + } + + const GraphNode * transpose = graph.index().producer(x_transposed->id); + if (transpose == nullptr || transpose->op != GGML_OP_TRANSPOSE || transpose->inputs.size() != 1) { + return {}; + } + const Value * x_layout = graph_value(graph, transpose->inputs[0]); + const Value * x = root_alias_value(graph, x_layout); + if (!is_f32(x_layout) || !x_layout->contiguous || !is_f32(x) || !x->contiguous) { + return {}; + } + + const int64_t hidden_size = x_layout->ne[0]; + const int64_t token_count = x_layout->ne[1]; + const int64_t sequence_count = x_layout->ne[2]; + if (!supported_hidden_size(hidden_size) || token_count < 1 || token_count > 512 || sequence_count < 1 || + sequence_count > 3 || !is_shape(*x_layout, hidden_size, token_count, sequence_count, 1) || + !is_shape(*x_transposed, token_count, hidden_size, sequence_count, 1) || + !is_shape(*state, 3, hidden_size, sequence_count, 1) || + !is_shape(*window, token_count + 3, hidden_size, sequence_count, 1)) { + return {}; + } + if (x_layout->nb[0] != sizeof(float) || x_layout->nb[1] != static_cast(hidden_size) * sizeof(float) || + x_layout->nb[2] != static_cast(hidden_size * token_count) * sizeof(float) || + x_transposed->nb[0] != static_cast(hidden_size) * sizeof(float) || + x_transposed->nb[1] != sizeof(float) || + x_transposed->nb[2] != static_cast(hidden_size * token_count) * sizeof(float) || + state->nb[0] != sizeof(float) || state->nb[1] != 3 * sizeof(float) || + state->nb[2] != static_cast(3 * hidden_size) * sizeof(float) || window->nb[0] != sizeof(float) || + window->nb[1] != static_cast(token_count + 3) * sizeof(float) || + window->nb[2] != static_cast((token_count + 3) * hidden_size) * sizeof(float)) { + return {}; + } + + const std::vector & window_consumers = graph.index().consumers(window->id); + if (window_consumers.size() < 2 || window_consumers.size() > 6) { + return {}; + } + std::vector state_tail_views; + const GraphNode * ssm = nullptr; + for (const GraphNode * consumer : window_consumers) { + if (consumer != nullptr && consumer->op == GGML_OP_VIEW) { + state_tail_views.push_back(consumer); + } else if (consumer != nullptr && consumer->op == GGML_OP_SSM_CONV) { + if (ssm != nullptr) { + return {}; + } + ssm = consumer; + } else { + return {}; + } + } + if (state_tail_views.empty() || state_tail_views.size() > 5 || ssm == nullptr || ssm->inputs.size() != 2 || + ssm->inputs[0] != window->id) { + return {}; + } + + const Value * filter = graph_value(graph, ssm->inputs[1]); + const Value * ssm_output = graph_value(graph, ssm->output); + if (!is_f32(filter) || !is_f32(ssm_output) || !filter->contiguous || !ssm_output->contiguous || + !is_shape(*filter, 4, hidden_size, 1, 1) || + !is_shape(*ssm_output, hidden_size, token_count, sequence_count, 1) || filter->nb[0] != sizeof(float) || + filter->nb[1] != 4 * sizeof(float)) { + return {}; + } + + std::vector cache_updates; + for (const GraphNode * state_tail_view : state_tail_views) { + if (state_tail_view == nullptr || state_tail_view->inputs.size() != 1) { + return {}; + } + const Value * state_tail = graph_value(graph, state_tail_view->output); + if (!is_f32(state_tail) || state_tail->alias_source != window->id || + state_tail->storage_offset < window->storage_offset || + (state_tail->storage_offset - window->storage_offset) % sizeof(float) != 0 || + !is_shape(*state_tail, 3, hidden_size, sequence_count, 1)) { + return {}; + } + const int64_t source_row = + static_cast((state_tail->storage_offset - window->storage_offset) / sizeof(float)); + if (source_row < 0 || source_row > token_count) { + return {}; + } + const std::vector & tail_consumers = graph.index().consumers(state_tail->id); + if (tail_consumers.size() != 1 || tail_consumers.front() == nullptr || + tail_consumers.front()->op != GGML_OP_CPY || tail_consumers.front()->inputs.size() != 2 || + tail_consumers.front()->inputs[0] != state_tail->id) { + return {}; + } + const GraphNode * cache_copy = tail_consumers.front(); + const Value * cache_target = graph_value(graph, cache_copy->inputs[1]); + const Value * cache = graph_value(graph, cache_copy->output); + const GraphNode * cache_view = cache_target != nullptr ? graph.index().producer(cache_target->id) : nullptr; + if (!is_f32(cache_target) || !is_f32(cache) || cache_view == nullptr || cache_view->op != GGML_OP_VIEW || + cache_view->inputs.size() != 1 || !same_full_value_range(*cache_target, *cache) || + cache->byte_count != static_cast(3 * hidden_size * sequence_count) * sizeof(float)) { + return {}; + } + cache_updates.push_back({ state_tail_view, cache_view, cache_copy, state_tail, cache, source_row }); + } + + std::sort(cache_updates.begin(), cache_updates.end(), + [](const auto & lhs, const auto & rhs) { return lhs.cache->storage_offset < rhs.cache->storage_offset; }); + for (size_t slot = 0; slot < cache_updates.size(); ++slot) { + const int64_t expected_source_row = std::max(0, token_count - static_cast(slot)); + if (cache_updates[slot].source_row != expected_source_row || + cache_updates[slot].cache->storage != cache_updates.front().cache->storage) { + return {}; + } + for (size_t prior = 0; prior < slot; ++prior) { + if (ranges_overlap(*cache_updates[slot].cache, *cache_updates[prior].cache)) { + return {}; + } + } + } + + const std::vector & ssm_consumers = graph.index().consumers(ssm_output->id); + if (ssm_consumers.size() != 1 || ssm_consumers.front() == nullptr || ssm_consumers.front()->op != GGML_OP_UNARY || + ssm_consumers.front()->inputs.size() != 1) { + return {}; + } + const GraphNode * silu = ssm_consumers.front(); + const UnaryParams * unary = op_params_as(silu->params); + const Value * output = graph_value(graph, silu->output); + if (unary == nullptr || unary->op != UnaryKind::Silu || !is_f32(output) || !output->contiguous || + !is_shape(*output, hidden_size, token_count, sequence_count, 1)) { + return {}; + } + + if (!distinct_storage(*state, *x) || !distinct_storage(*state, *filter) || !distinct_storage(*x, *filter) || + (!distinct_storage(*output, *x) && !same_full_value_range(*output, *x)) || !distinct_storage(*output, *state) || + !distinct_storage(*output, *filter)) { + return {}; + } + for (const SsmConvCacheUpdate & update : cache_updates) { + if (!distinct_storage(*state, *update.cache) || !distinct_storage(*x, *update.cache) || + !distinct_storage(*filter, *update.cache) || !distinct_storage(*output, *update.cache)) { + return {}; + } + } + + match.concat = node; + match.ssm = ssm; + match.silu = silu; + match.state = state; + match.x = x; + match.filter = filter; + match.output = output; + match.cache_updates = std::move(cache_updates); + match.hidden_size = hidden_size; + match.token_count = token_count; + match.sequence_count = sequence_count; + return match; +} + +static SsmConvCoreMatch match_ssm_conv_core(const Graph & graph, const GraphNode * node) { + SsmConvCoreMatch match; + if (node == nullptr || node->op != GGML_OP_SSM_CONV || node->inputs.size() != 2) { + return match; + } + + const Value * window = graph_value(graph, node->inputs[0]); + const Value * filter = graph_value(graph, node->inputs[1]); + const Value * conv_output = graph_value(graph, node->output); + if (!is_f32(window) || !is_f32(filter) || !is_f32(conv_output) || !packed_f32_layout(*window) || + !packed_f32_layout(*filter) || !packed_f32_layout(*conv_output)) { + return {}; + } + + const int64_t d_conv = filter->ne[0]; + const int64_t d_inner = filter->ne[1]; + const int64_t token_count = window->ne[0] - d_conv + 1; + const int64_t sequence_count = window->ne[2]; + if (d_conv < 1 || d_conv > 16 || !supported_hidden_size(d_inner) || token_count < 1 || token_count > 512 || + sequence_count < 1 || sequence_count > 4 || window->ne[1] != d_inner || window->ne[3] != 1 || + filter->ne[2] != 1 || filter->ne[3] != 1 || !is_shape(*conv_output, d_inner, token_count, sequence_count, 1) || + ranges_overlap(*window, *conv_output) || ranges_overlap(*filter, *conv_output)) { + return {}; + } + + match.ssm = node; + match.window = window; + match.filter = filter; + match.conv_output = conv_output; + match.output = conv_output; + match.d_conv = d_conv; + match.d_inner = d_inner; + match.token_count = token_count; + match.sequence_count = sequence_count; + return match; +} + +static bool try_match_unary_fusion(const Graph & graph, SsmConvCoreMatch & match) { + if (!graph.has_index()) { + return false; + } + const std::vector & consumers = graph.index().consumers(match.conv_output->id); + if (consumers.size() != 1 || consumers.front() == nullptr || consumers.front()->op != GGML_OP_UNARY || + consumers.front()->inputs.size() != 1) { + return false; + } + + const GraphNode * unary = consumers.front(); + const UnaryParams * params = op_params_as(unary->params); + const Value * output = graph_value(graph, unary->output); + if (params == nullptr || !unary_kind_supported(params->op) || output == nullptr || output->type != GGML_TYPE_F32 || + !packed_f32_layout(*output) || !same_shape(*output, *match.conv_output) || + ranges_overlap(*match.window, *output) || ranges_overlap(*match.filter, *output)) { + return false; + } + + match.unary = unary; + match.output = output; + match.unary_op = params->op; + return true; +} + +static bool try_match_binary_fusion(const Graph & graph, SsmConvCoreMatch & match) { + if (!graph.has_index()) { + return false; + } + const std::vector & consumers = graph.index().consumers(match.conv_output->id); + if (consumers.size() != 1 || consumers.front() == nullptr || consumers.front()->inputs.size() != 2) { + return false; + } + + const GraphNode * binary = consumers.front(); + const CommonBinaryMatch binary_match = + common_match_binary_consumer(binary, match.conv_output->id, kCommonArithmeticBinaryKinds); + if (!binary_match.matched()) { + return false; + } + + const Value * operand = graph_value(graph, binary_match.operand); + const Value * output = graph_value(graph, binary->output); + if (!is_f32(operand) || output == nullptr || output->type != GGML_TYPE_F32 || !packed_f32_layout(*operand) || + !packed_f32_layout(*output) || !same_shape(*operand, *match.conv_output) || + !same_shape(*output, *match.conv_output) || ranges_overlap(*match.window, *output) || + ranges_overlap(*match.filter, *output) || ranges_overlap(*operand, *output)) { + return false; + } + + match.binary = binary; + match.binary_operand = operand; + match.output = output; + match.binary_op = binary_match.kind; + match.binary_lhs = binary_match.producer_is_lhs; + return true; +} + +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); + 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); +} + +static bool append_ssm_conv_covered_nodes(const DispatchMatchContext & context, + const SsmConvPrefillMatch & match, + DispatchMatch & dispatch_match) { + if (!append_covered_node(context, match.concat, dispatch_match) || + !append_covered_node(context, match.ssm, dispatch_match) || + !append_covered_node(context, match.silu, dispatch_match)) { + return false; + } + for (const SsmConvCacheUpdate & update : match.cache_updates) { + if (!append_covered_node(context, update.state_tail_view, dispatch_match) || + !append_covered_node(context, update.cache_view, dispatch_match) || + !append_covered_node(context, update.cache_copy, dispatch_match)) { + return false; + } + } + return true; +} + +static bool match_ssm_conv_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const SsmConvPrefillMatch match = match_ssm_conv_prefill(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + if (!append_ssm_conv_covered_nodes(context, match, dispatch_match)) { + return false; + } + + const CommandPlanGeneratedResource * edges = + context.plan.metadata.find_generated_resource(match.x->id, GeneratedResourceRole::Conv4Edges); + if (edges != nullptr) { + Dispatch finish; + finish.kernel = make_kernel_specialization(kSsmConvFinishKernel); + set_compile_parameter(finish.kernel, "llm.ssm_conv.generic.d_inner", match.hidden_size); + finish.bindings.push_back({ match.state->id, 0, match.state->byte_count }); + finish.bindings.push_back({ match.filter->id, 0, match.filter->byte_count }); + finish.bindings.push_back({ edges->generated_value, 0, edges->byte_count }); + finish.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + const Value * cache = match.cache_updates.front().cache; + finish.bindings.push_back({ cache->id, 0, cache->byte_count }); + dispatch_match.dispatches.push_back(std::move(finish)); + return true; + } + + const bool optimized_prefill = match.token_count == 512 && match.sequence_count == 1 && + match.cache_updates.size() == 1 && match.hidden_size >= 8192 && + match.hidden_size <= 10240; + if (!optimized_prefill) { + Dispatch ssm_dispatch; + const bool decode = match.token_count == 1 && match.sequence_count == 1 && match.cache_updates.size() == 1; + if (decode) { + ssm_dispatch.kernel = make_kernel_specialization(kSsmConvDecodeKernel); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.decode.d_inner", match.hidden_size); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.decode.workgroup_size", 256); + } else { + ssm_dispatch.kernel = make_kernel_specialization(kSsmConvRollbackKernel); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.rollback.d_inner", match.hidden_size); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.rollback.n_t", match.token_count); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.rollback.n_s", match.sequence_count); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.rollback.cache_count", match.cache_updates.size()); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.rollback.workgroup_size", 256); + } + ssm_dispatch.bindings.push_back({ match.state->id, 0, match.state->byte_count }); + ssm_dispatch.bindings.push_back({ match.x->id, 0, match.x->byte_count }); + ssm_dispatch.bindings.push_back({ match.filter->id, 0, match.filter->byte_count }); + ssm_dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + if (decode) { + const Value * cache = match.cache_updates.front().cache; + ssm_dispatch.bindings.push_back({ cache->id, 0, cache->byte_count }); + } else { + for (size_t slot = 0; slot < 5; ++slot) { + const Value * cache = match.cache_updates[std::min(slot, match.cache_updates.size() - 1)].cache; + ssm_dispatch.bindings.push_back({ cache->id, 0, cache->byte_count }); + } + } + dispatch_match.dispatches.push_back(std::move(ssm_dispatch)); + return true; + } + + const size_t snapshot_bytes = static_cast(64 * match.hidden_size) * sizeof(float); + const ValueId snapshot = context.next_plan_value; + dispatch_match.transients.push_back({ snapshot, "llm.ssm_conv.window_snapshot", snapshot_bytes, 256 }); + + Dispatch snapshot_dispatch; + snapshot_dispatch.kernel = make_kernel_specialization(kSsmConvSnapshotKernel); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.d_conv", 4); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.d_inner", match.hidden_size); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.n_t", match.token_count); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.n_s", match.sequence_count); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.state_row_stride", 1); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.state_channel_stride", 3); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.x_row_stride", match.hidden_size); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.dst_row_stride", match.hidden_size); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.cache_row_stride", 1); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.cache_channel_stride", 3); + set_compile_parameter(snapshot_dispatch.kernel, "llm.ssm_conv.snapshot.workgroup_size", 256); + snapshot_dispatch.bindings.push_back({ match.state->id, 0, match.state->byte_count }); + snapshot_dispatch.bindings.push_back({ match.x->id, 0, match.x->byte_count }); + snapshot_dispatch.bindings.push_back({ snapshot, 0, snapshot_bytes }); + const Value * cache = match.cache_updates.front().cache; + snapshot_dispatch.bindings.push_back({ cache->id, 0, cache->byte_count }); + + Dispatch ssm_dispatch; + ssm_dispatch.kernel = make_kernel_specialization(kSsmConvPrefillKernel); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.prefill.d_conv", 4); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.prefill.d_inner", match.hidden_size); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.prefill.n_t", match.token_count); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.prefill.n_s", match.sequence_count); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.prefill.state_row_stride", match.hidden_size); + set_compile_parameter(ssm_dispatch.kernel, "llm.ssm_conv.prefill.x_row_stride", match.hidden_size); + ssm_dispatch.bindings.push_back({ snapshot, 0, snapshot_bytes }); + ssm_dispatch.bindings.push_back({ match.x->id, 0, match.x->byte_count }); + ssm_dispatch.bindings.push_back({ match.filter->id, 0, match.filter->byte_count }); + ssm_dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + + dispatch_match.dispatches.push_back(std::move(snapshot_dispatch)); + dispatch_match.dispatches.push_back(std::move(ssm_dispatch)); + return true; +} + +static bool match_ssm_conv_generic_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + SsmConvCoreMatch match = match_ssm_conv_core(context.graph, context.root_node); + if (!match.matched()) { + return false; + } + + try_match_binary_fusion(context.graph, match) || try_match_unary_fusion(context.graph, match); + + if (!append_covered_node(context, match.ssm, dispatch_match)) { + return false; + } + if (match.has_binary_fusion() && !append_covered_node(context, match.binary, dispatch_match)) { + return false; + } + if (match.has_unary_fusion() && !append_covered_node(context, match.unary, dispatch_match)) { + return false; + } + + Dispatch dispatch; + dispatch.kernel = + make_kernel_specialization(match.has_binary_fusion() ? kSsmConvGenericBinaryKernel : kSsmConvGenericKernel); + set_generic_ssm_conv_parameters(dispatch.kernel, match); + dispatch.bindings.push_back({ match.window->id, 0, match.window->byte_count }); + dispatch.bindings.push_back({ match.filter->id, 0, match.filter->byte_count }); + if (match.has_binary_fusion()) { + common_set_binary_compile_parameters(dispatch.kernel, "llm.ssm_conv.generic.binary_op", + "llm.ssm_conv.generic.binary_lhs", match.binary_op, match.binary_lhs); + dispatch.bindings.push_back({ match.binary_operand->id, 0, match.binary_operand->byte_count }); + } + dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { + const CommonMulMatMatch mat = + common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatConv4Kernel, false); + if (!mat.matched() || !context.graph.has_index() || mat.token_count != 512 || + (mat.weight->type != GGML_TYPE_Q4_K && mat.weight->type != GGML_TYPE_Q6_K) || + mat.weight->alias_source.value >= 0 || mat.input_size % 256 != 0 || mat.output_size % 64 != 0 || + mat.output_size / 64 < 32 || + !common_mul_mat_uses_k16_major_f16(mat.weight_format, mat.input_size, mat.output_size, mat.token_count)) { + return false; + } + + const Value * value = mat.output; + std::vector layouts; + const GraphNode * concat = nullptr; + while (value != nullptr && value->kind == ValueKind::Transient) { + const auto & consumers = context.graph.index().consumers(value->id); + if (consumers.size() != 1 || consumers.front() == nullptr) { + return false; + } + const GraphNode * consumer = consumers.front(); + if (consumer->op == GGML_OP_CONCAT) { + concat = consumer; + break; + } + if (!is_layout_alias_node(context.graph, *consumer)) { + return false; + } + layouts.push_back(consumer); + value = graph_value(context.graph, consumer->output); + } + const SsmConvPrefillMatch conv = match_ssm_conv_prefill(context.graph, concat); + if (!conv.matched() || conv.x->id != mat.output->id || conv.sequence_count != 1 || conv.token_count != 512 || + conv.cache_updates.size() != 1) { + return false; + } + const Value * window = graph_value(context.graph, conv.concat->output); + const Value * intermediate = graph_value(context.graph, conv.ssm->output); + if (window->kind != ValueKind::Transient || intermediate->kind != ValueKind::Transient) { + return false; + } + bool deferred_state = false; + for (const Value * input : {conv.state, conv.filter}) { + const GraphNode * producer = context.graph.index().producer(input->id); + size_t index = 0; + if (producer != nullptr && + (!context.graph.index().node_index(producer, index) || !context.covered_nodes[index])) { + if (input == conv.filter) { + return false; + } + deferred_state = true; + } + } + if (deferred_state && (conv.hidden_size < 8192 || conv.hidden_size > 10240)) { + return false; + } + const Value * cache = conv.cache_updates.front().cache; + for (const Value * output : {conv.output, cache}) { + if (!distinct_storage(*output, *mat.input) || !distinct_storage(*output, *mat.weight)) { + return false; + } + } + if (!append_covered_node(context, context.root_node, dispatch_match) || + (!deferred_state && !append_ssm_conv_covered_nodes(context, conv, dispatch_match))) { + return false; + } + for (const GraphNode * layout : layouts) { + if (!append_covered_node(context, layout, dispatch_match)) { + return false; + } + } + + DispatchBinding activation; + if (!common_prepare_k16_major_f16_input(context, *mat.input, mat.input_size, mat.token_count, dispatch_match, + activation)) { + return false; + } + Dispatch dispatch; + dispatch.kernel = make_kernel_specialization(deferred_state ? kMulMatConv4InteriorKernel : kMulMatConv4Kernel); + dispatch.kernel.integer_parameters.emplace("token_count", mat.token_count); + set_compile_parameter(dispatch.kernel, "ggml.mul_mat.input_size", mat.input_size); + set_compile_parameter(dispatch.kernel, "ggml.mul_mat.output_size", mat.output_size); + set_compile_parameter(dispatch.kernel, "ggml.mul_mat.weight_format", mat.weight->type == GGML_TYPE_Q4_K ? 4 : 6); + dispatch.bindings.push_back(activation); + if (mat.weight->type == GGML_TYPE_Q4_K) { + dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count, kQ4KPackedK256Row64Layout, + mat.weight->type, mat.input_size, mat.output_size, mat.weight->byte_count}); + } else { + dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count}); + } + if (!deferred_state) { + dispatch.bindings.push_back({conv.state->id, 0, conv.state->byte_count}); + } + dispatch.bindings.push_back({conv.filter->id, 0, conv.filter->byte_count}); + dispatch.bindings.push_back({conv.output->id, 0, conv.output->byte_count}); + if (deferred_state) { + const size_t edge_bytes = static_cast(6 * conv.hidden_size) * sizeof(float); + const ValueId edges(context.next_plan_value.value + static_cast(dispatch_match.transients.size())); + dispatch_match.transients.push_back({ edges, "llm.ssm_conv.projection_edges", edge_bytes, 256 }); + Status status; + if (!dispatch_match.metadata.append_generated_resource( + { mat.output->id, GeneratedResourceRole::Conv4Edges, edges, edge_bytes, {} }, status)) { + dispatch_match.status.append(status); + return false; + } + dispatch.bindings.push_back({ edges, 0, edge_bytes }); + } else { + dispatch.bindings.push_back({cache->id, 0, cache->byte_count}); + } + dispatch_match.dispatches.push_back(std::move(dispatch)); + return true; +} + +} // namespace + +void register_llm_ssm_conv_dispatch(DispatchRegistryBuilder & registry) { + registry.add({ + "llm.ssm_conv.quantized_prefill", + GGML_OP_MUL_MAT, + DispatchMatchKind::Fused, + 306, + DispatchSource::Llm, + match_mul_mat_conv4_dispatch, + }); + registry.add({ + "llm.ssm_conv.dconv4_silu", + GGML_OP_CONCAT, + DispatchMatchKind::Fused, + 200, + DispatchSource::Llm, + match_ssm_conv_prefill_dispatch, + }); + registry.add({ + "llm.ssm_conv.lfm_dconv3", + GGML_OP_CONCAT, + DispatchMatchKind::Fused, + 201, + DispatchSource::Llm, + match_lfm_dconv3_dispatch, + }); + registry.add({ + "llm.ssm_conv.generic_f32", + GGML_OP_SSM_CONV, + DispatchMatchKind::Fused, + 50, + DispatchSource::Llm, + match_ssm_conv_generic_dispatch, + }); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.h b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.h new file mode 100644 index 000000000000..a17e37af73d7 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_llm_ssm_conv_dispatch(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-profiles.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-profiles.h similarity index 100% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-profiles.h rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-profiles.h diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-shapes.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-shapes.h similarity index 100% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-llm-shapes.h rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-shapes.h diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp similarity index 92% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.cpp rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp index 3b88b7d26ca6..890159cd6f09 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.cpp @@ -21,10 +21,10 @@ static constexpr KernelCatalogRef kQwenRouterProjectionTop8FusedDecodeF32Kernel GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_router_projection_top8_fused_decode_f32"); static constexpr KernelCatalogRef kQwenRouterProjectionF32FourRowWave32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_router_projection_f32_four_row_wave32"); -static constexpr KernelCatalogRef kQwenBuildExpertTableKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_build_expert_table"); -static constexpr KernelCatalogRef kQwenBuildExpertPartitionTableKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_build_expert_partition_table"); +static constexpr KernelCatalogRef kMoeBuildExpertTableKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_moe_build_expert_table"); +static constexpr KernelCatalogRef kMoeBuildExpertPartitionTableKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_moe_build_expert_partition_table"); static constexpr KernelCatalogRef kQwenBuildExpertTablePartitionPrefill512Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_build_expert_table_partition_prefill_512"); @@ -361,6 +361,19 @@ static void add_routed_gate_up_compile_parameters(Dispatch & dispatch, const Rou to_config_value(router_match.route_count)); } +static void add_moe_routing_compile_parameters(Dispatch & dispatch, const RouterTop8Match & router_match) { + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", + to_config_value(router_match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.route_count", + to_config_value(router_match.route_count)); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.expert_count", + to_config_value(router_match.expert_count)); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.descriptor_expert_mask", "127"); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.descriptor_partition_shift", "7"); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.descriptor_row_count_shift", "13"); + dispatch.kernel.compile_parameters.emplace("ggml.moe_routing.partition_workgroup_size", "128"); +} + struct RouterProjectionTop8Match { const GraphNode * projection = nullptr; const GraphNode * softmax = nullptr; @@ -451,6 +464,25 @@ static RouterProjectionMatch match_moe_router_projection_f32(const DispatchMatch return match; } +static bool has_common_matmul_fusible_unary_consumer(const DispatchMatchContext & context, const Value & output) { + const std::vector & consumers = context.graph.index().consumers(output.id); + if (consumers.size() != 1 || consumers.front() == nullptr) { + return false; + } + + const GraphNode * unary = consumers.front(); + size_t unary_index = 0; + if (!context.graph.index().node_index(unary, unary_index) || unary_index >= context.covered_nodes.size() || + context.covered_nodes[unary_index] || unary->inputs.size() != 1 || unary->inputs[0] != output.id) { + return false; + } + + const UnaryParams * params = op_params_as(unary->params); + const Value * result = graph_value(context.graph, unary->output); + return params != nullptr && unary_kind_supported(params->op) && result != nullptr && + result->type == GGML_TYPE_F32 && result->contiguous && same_shape(output, *result); +} + } // namespace static bool match_moe_router_projection_f32_dispatch(const DispatchMatchContext & context, @@ -459,6 +491,9 @@ static bool match_moe_router_projection_f32_dispatch(const DispatchMatchContext if (!match.matched()) { return false; } + if (has_common_matmul_fusible_unary_consumer(context, *match.output)) { + return false; + } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(kQwenRouterProjectionF32FourRowWave32Kernel); @@ -623,26 +658,22 @@ static bool match_moe_router_top8_dispatch(const DispatchMatchContext & context, dispatch_match.dispatches.push_back(std::move(expert_table_partition_dispatch)); } else { Dispatch expert_table_dispatch; - expert_table_dispatch.kernel = make_kernel_specialization(kQwenBuildExpertTableKernel); + expert_table_dispatch.kernel = make_kernel_specialization(kMoeBuildExpertTableKernel); expert_table_dispatch.kernel.integer_parameters.emplace("token_count", router_match.token_count); expert_table_dispatch.kernel.integer_parameters.emplace("route_count", router_match.route_count); expert_table_dispatch.kernel.integer_parameters.emplace("route_stride", router_match.route_stride); expert_table_dispatch.kernel.integer_parameters.emplace("expert_count", router_match.expert_count); - expert_table_dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", - to_config_value(router_match.token_count)); - add_routed_gate_up_compile_parameters(expert_table_dispatch, router_match); + add_moe_routing_compile_parameters(expert_table_dispatch, router_match); expert_table_dispatch.bindings.push_back({ router_match.route_ids->id, 0, route_id_length }); expert_table_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); dispatch_match.dispatches.push_back(std::move(expert_table_dispatch)); Dispatch partition_table_dispatch; - partition_table_dispatch.kernel = make_kernel_specialization(kQwenBuildExpertPartitionTableKernel); + partition_table_dispatch.kernel = make_kernel_specialization(kMoeBuildExpertPartitionTableKernel); partition_table_dispatch.kernel.integer_parameters.emplace("token_count", router_match.token_count); partition_table_dispatch.kernel.integer_parameters.emplace("route_count", router_match.route_count); partition_table_dispatch.kernel.integer_parameters.emplace("expert_count", router_match.expert_count); - partition_table_dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", - to_config_value(router_match.token_count)); - add_routed_gate_up_compile_parameters(partition_table_dispatch, router_match); + add_moe_routing_compile_parameters(partition_table_dispatch, router_match); partition_table_dispatch.bindings.push_back({ expert_table_value, 0, expert_table_bytes }); partition_table_dispatch.bindings.push_back({ partition_table_value, 0, partition_table_bytes }); dispatch_match.dispatches.push_back(std::move(partition_table_dispatch)); @@ -670,7 +701,7 @@ void register_moe_router_dispatches(DispatchRegistryBuilder & registry) { registry.add({ "llm.moe_router.projection_f32_four_row_wave32", GGML_OP_MUL_MAT, - DispatchMatchKind::SingleOp, + DispatchMatchKind::Fused, 90, DispatchSource::Llm, match_moe_router_projection_f32_dispatch, diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.h similarity index 80% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.h rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.h index 296282072a7f..ff174cff5d12 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-moe-router.h +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-moe-router.h @@ -1,6 +1,6 @@ #pragma once -#include "dispatch-registry.h" +#include "../dispatch-registry.h" namespace ggml::hrx { diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp similarity index 94% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.cpp rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp index 28a34bc01836..9b0dfd1cff88 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.cpp @@ -1,5 +1,6 @@ #include "dispatch-qwen-attention-postprocess.h" +#include "../common/dispatch-rope-utils.h" #include "dispatch-llm-shapes.h" #include "ggml.h" #include "graph/graph-matcher.h" @@ -21,7 +22,7 @@ static constexpr KernelCatalogRef kQwenAttentionPostprocessF32F16Kernel = static constexpr KernelCatalogRef kQwenAttentionQkvPostprocessFusedDecodeKernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_attention_qkv_postprocess_fused_decode"); static constexpr KernelCatalogRef kQwenAttentionContextBaseCaptureKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen_attention_context_base_capture"); + GGML_HRX_KERNEL_REF("qwen", "qwen_attention_context_base_capture"); static constexpr KernelCatalogRef kQwenAttentionMetadataKernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen_attention_metadata"); static constexpr int64_t kQwenAttentionHeadSize = 128; @@ -60,22 +61,21 @@ static bool is_qwen_rms_norm_epsilon(float eps) { static bool is_qwen_implicit_rope_contract(const RopeParams & params) { return params.n_dims == kQwenAttentionHeadSize && params.mode == GGML_ROPE_TYPE_NEOX && std::isfinite(params.freq_base) && params.freq_base > 0.0f && std::isfinite(params.freq_scale) && - params.freq_scale > 0.0f && params.ext_factor == 0.0f && params.attn_factor == 1.0f; + params.freq_scale > 0.0f && std::isfinite(params.ext_factor) && std::isfinite(params.attn_factor); } -static bool build_inverse_frequency_table(const GraphNode & rope, std::vector & data) { +static bool build_inverse_frequency_table(const GraphNode & rope, std::vector & data, float & rope_mscale) { const RopeParams * params = op_params_as(rope.params); if (params == nullptr || !is_qwen_implicit_rope_contract(*params)) { return false; } - data.resize(static_cast(params->n_dims / 2) * sizeof(float)); - const float theta_scale = std::pow(params->freq_base, -2.0f / static_cast(params->n_dims)); - float theta = params->freq_scale; - for (int i = 0; i < params->n_dims / 2; ++i) { - std::memcpy(data.data() + static_cast(i) * sizeof(float), &theta, sizeof(theta)); - theta *= theta_scale; + RopeFrequencyTable table; + if (!build_rope_frequency_table(*params, params->n_dims, table)) { + return false; } + data = std::move(table.data); + rope_mscale = table.mscale; return true; } @@ -209,6 +209,7 @@ struct NormRopeChain { const Value * inverse_freqs = nullptr; size_t inverse_freqs_byte_count = 0; std::vector inverse_freqs_data; + float rope_mscale = 1.0f; const Value * output = nullptr; int64_t token_count = 0; int64_t head_count = 0; @@ -284,9 +285,9 @@ struct AttentionPostprocessMatch { static bool matching_inverse_frequencies(const NormRopeChain & lhs, const NormRopeChain & rhs) { if (lhs.inverse_freqs != nullptr || rhs.inverse_freqs != nullptr) { return lhs.inverse_freqs != nullptr && rhs.inverse_freqs != nullptr && - lhs.inverse_freqs->id == rhs.inverse_freqs->id; + lhs.inverse_freqs->id == rhs.inverse_freqs->id && lhs.rope_mscale == rhs.rope_mscale; } - return lhs.inverse_freqs_data == rhs.inverse_freqs_data; + return lhs.inverse_freqs_data == rhs.inverse_freqs_data && lhs.rope_mscale == rhs.rope_mscale; } static bool has_qwen_rope_params(const GraphNode & node) { @@ -401,14 +402,27 @@ static bool match_norm_rope_chain_from_reshape(const Graph & graph, size_t inverse_freqs_byte_count = 0; const Value * inverse_freqs = nullptr; std::vector inverse_freqs_data; + float rope_mscale = 1.0f; if (rope->inputs.size() == 3) { inverse_freqs = graph_value(graph, rope->inputs[2]); if (inverse_freqs == nullptr || !is_inverse_frequency_table(*inverse_freqs)) { log_attention_reject(status, graph, reshape, label + " explicit inverse-frequency table is incompatible"); return false; } + const RopeParams * params = op_params_as(rope->params); + if (params == nullptr || !is_qwen_implicit_rope_contract(*params)) { + log_attention_reject(status, graph, reshape, + label + " explicit inverse-frequency table has unsupported ROPE params=" + + rope_params_summary(*rope)); + return false; + } + RopeFrequencyTable table; + if (!build_rope_frequency_table(*params, params->n_dims, table)) { + return false; + } + rope_mscale = table.mscale; inverse_freqs_byte_count = inverse_freqs->byte_count; - } else if (build_inverse_frequency_table(*rope, inverse_freqs_data)) { + } else if (build_inverse_frequency_table(*rope, inverse_freqs_data, rope_mscale)) { inverse_freqs_byte_count = inverse_freqs_data.size(); } else { log_attention_reject(status, graph, reshape, @@ -431,6 +445,7 @@ static bool match_norm_rope_chain_from_reshape(const Graph & graph, chain.inverse_freqs = inverse_freqs; chain.inverse_freqs_byte_count = inverse_freqs_byte_count; chain.inverse_freqs_data = std::move(inverse_freqs_data); + chain.rope_mscale = rope_mscale; chain.output = output; return true; } @@ -744,6 +759,20 @@ static bool append_attention_metadata_initialization(const DispatchMatchContext if (mask == nullptr || mask->type != GGML_TYPE_F16 || mask->ne[0] <= 0 || mask->ne[1] != match.query.token_count) { return false; } + int64_t context_capacity = mask->ne[0]; + size_t mask_byte_count = mask->byte_count; + if (context_capacity > 32768 && match.query.token_count <= 32768) { + context_capacity = match.query.token_count; + mask_byte_count = static_cast(match.query.token_count) * static_cast(match.query.token_count) * + sizeof(ggml_fp16_t); + Status metadata_status; + if (!dispatch_match.metadata.append_alternate_value( + { mask->id, mask->id, GGML_TYPE_F16, mask_byte_count, "qwen.attention.compact_mask" }, + metadata_status)) { + dispatch_match.status.append(metadata_status); + return false; + } + } const ValueId control = next_match_transient_value(context, dispatch_match); dispatch_match.transients.push_back({ control, "qwen.attention.control", sizeof(int32_t), 16 }); @@ -757,12 +786,12 @@ static bool append_attention_metadata_initialization(const DispatchMatchContext Dispatch metadata; metadata.kernel = make_kernel_specialization(kQwenAttentionMetadataKernel); metadata.kernel.integer_parameters.emplace("token_count", match.query.token_count); - metadata.kernel.integer_parameters.emplace("context_capacity", mask->ne[0]); + metadata.kernel.integer_parameters.emplace("context_capacity", context_capacity); metadata.bindings.push_back({ control, 0, sizeof(int32_t) }); metadata.bindings.push_back({ match.query.positions->id, 0, match.query.positions->byte_count }); metadata.bindings.push_back({ match.key.cache_indices->id, 0, match.key.cache_indices->byte_count }); metadata.bindings.push_back({ match.value.cache_indices->id, 0, match.value.cache_indices->byte_count }); - metadata.bindings.push_back({ mask->id, 0, mask->byte_count }); + metadata.bindings.push_back({ mask->id, 0, mask_byte_count }); dispatch_match.initialization_dispatches.push_back(std::move(metadata)); return true; } @@ -871,6 +900,8 @@ static bool match_qwen_attention_qkv_postprocess_fused_decode_dispatch(const Dis value_weight->type == GGML_TYPE_Q6_K ? "1" : "0"); dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.head_size", to_config_value(kQwenAttentionHeadSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.rope_mscale", + rope_mscale_config_value(match.query.rope_mscale)); dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.query.token_count)); @@ -932,6 +963,8 @@ static bool match_qwen_attention_postprocess_dispatch(const DispatchMatchContext to_config_value(match.query.head_count * kQwenAttentionHeadSize)); dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.key_value_size", to_config_value(match.key.key.head_count * kQwenAttentionHeadSize)); + dispatch.kernel.compile_parameters.emplace("qwen3_moe.attention.rope_mscale", + rope_mscale_config_value(match.query.rope_mscale)); dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.query.token_count)); dispatch.bindings.push_back({ match.query.positions->id, 0, match.query.positions->byte_count }); diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.h similarity index 81% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.h rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.h index 0c3a8b6be354..66cb74342732 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-attention-postprocess.h +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-attention-postprocess.h @@ -1,6 +1,6 @@ #pragma once -#include "dispatch-registry.h" +#include "../dispatch-registry.h" namespace ggml::hrx { diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.cpp similarity index 65% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-qwen-matmul.cpp rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.cpp index 715751f5521a..fcd7597bb721 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/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" @@ -12,10 +14,6 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kQwenDenseLinearQ4KF16WmmaKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q4k_f16_wmma"); -static constexpr KernelCatalogRef kQwenDenseLinearQ6KF16WmmaKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q6k_f16_wmma"); static constexpr KernelCatalogRef kQwenDenseLinearQ4KQ8NextQ8Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"); static constexpr KernelCatalogRef kGgmlLinearQ6KQ8_1X4Kernel = @@ -40,10 +38,6 @@ static bool is_supported_dense_output_size(int64_t output_size) { return output_size >= 1 && output_size <= 262144; } -static bool is_qwen_endpoint_projection(int64_t input_size, int64_t output_size) { - return input_size == kQwenHiddenSize && output_size == kQwenVocabularyCount; -} - static std::string to_config_value(int64_t value) { return std::to_string(value); } @@ -58,7 +52,6 @@ struct QwenMatmulMatch { int64_t input_size = 0; int64_t output_size = 0; int64_t token_count = 0; - bool dense = false; bool matched() const { return input != nullptr && weight != nullptr && output != nullptr && kernel.id != kUncatalogedKernelId; @@ -90,23 +83,6 @@ struct QwenAttentionOutputNextQ8Match { } }; -struct QwenAttentionOutputAccumulateMatch { - const Value * input = nullptr; - const Value * weight = nullptr; - const Value * projection_output = nullptr; - const Value * residual_input = nullptr; - const Value * residual_output = nullptr; - const GraphNode * add_node = nullptr; - int64_t input_size = 0; - int64_t output_size = 0; - int64_t token_count = 0; - - bool matched() const { - return input != nullptr && weight != nullptr && projection_output != nullptr && residual_input != nullptr && - residual_output != nullptr && add_node != nullptr; - } -}; - static size_t q8_1_x4_byte_count(int64_t token_count, int64_t input_size) { if (token_count <= 0 || input_size <= 0) { return 0; @@ -222,43 +198,6 @@ static QwenMatmulMatch match_qwen_q6k_q8_matmul(const Graph & graph, const Graph return match; } -static QwenMatmulMatch match_qwen_decode_endpoint_q6k_matmul(const Graph & graph, const GraphNode * node) { - QwenMatmulMatch match; - if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2) { - return match; - } - - const Value * weight = graph_value(graph, node->inputs[0]); - const Value * input = graph_value(graph, node->inputs[1]); - const Value * output = graph_value(graph, node->output); - if (weight == nullptr || input == nullptr || output == nullptr || !is_2d(*weight) || !is_2d(*input) || - !is_2d(*output) || !weight->contiguous || !input->contiguous || !output->contiguous || - weight->type != GGML_TYPE_Q6_K || input->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32) { - 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 || - !is_qwen_decode_query_length(token_count) || !is_qwen_endpoint_projection(input_size, output_size) || - !is_supported_dense_input_size(input_size) || !is_supported_dense_output_size(output_size)) { - return {}; - } - - match.input = input; - match.weight = weight; - match.output = output; - match.input_value = input->id; - match.input_bytes = input->byte_count; - match.kernel = kQwenDenseLinearQ6KF16WmmaKernel; - match.input_size = input_size; - match.output_size = output_size; - match.token_count = token_count; - match.dense = true; - return match; -} - static QwenAttentionOutputNextQ8Match match_qwen_attention_output_next_q8(const DispatchMatchContext & context) { QwenAttentionOutputNextQ8Match match; const Graph & graph = context.graph; @@ -379,91 +318,36 @@ static QwenAttentionOutputNextQ8Match match_qwen_attention_output_next_q8(const return match; } -static QwenAttentionOutputAccumulateMatch match_qwen_attention_output_accumulate(const DispatchMatchContext & context) { - QwenAttentionOutputAccumulateMatch match; - const Graph & graph = context.graph; - const GraphNode * node = context.root_node; - if (node == nullptr || node->op != GGML_OP_MUL_MAT || node->inputs.size() != 2 || !graph.has_index()) { - return match; - } - - const Value * weight = graph_value(graph, node->inputs[0]); - const Value * input = graph_value(graph, node->inputs[1]); - const Value * projection_output = graph_value(graph, node->output); - if (weight == nullptr || input == nullptr || projection_output == nullptr || !is_2d(*weight) || !is_2d(*input) || - !is_2d(*projection_output)) { - return {}; - } - if (weight->type != GGML_TYPE_Q4_K || input->type != GGML_TYPE_F32 || projection_output->type != GGML_TYPE_F32 || - !weight->contiguous || !input->contiguous || !projection_output->contiguous) { - 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 (!is_qwen_prefill_query_length(token_count) || input->ne[0] != input_size || - projection_output->ne[0] != output_size || projection_output->ne[1] != token_count || input_size != 4096 || - output_size != kQwenHiddenSize) { - return {}; - } +} // namespace - const GraphNode * add_node = find_single_consumer_with_op(graph, projection_output->id, GGML_OP_ADD); - if (add_node == nullptr || add_node->inputs.size() != 2) { - return {}; - } - const Value * residual_input = nullptr; - if (add_node->inputs[0] == projection_output->id) { - residual_input = graph_value(graph, add_node->inputs[1]); - } else if (add_node->inputs[1] == projection_output->id) { - residual_input = graph_value(graph, add_node->inputs[0]); - } - const Value * residual_output = graph_value(graph, add_node->output); - if (residual_input == nullptr || residual_output == nullptr || residual_input->type != GGML_TYPE_F32 || - residual_output->type != GGML_TYPE_F32 || !same_value_layout(*projection_output, *residual_input) || - !same_value_layout(*projection_output, *residual_output) || - !value_has_no_uncovered_consumers_except(context, residual_input->id, add_node)) { - return {}; +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; } - - match.input = input; - match.weight = weight; - match.projection_output = projection_output; - match.residual_input = residual_input; - match.residual_output = residual_output; - match.add_node = add_node; - match.input_size = input_size; - match.output_size = output_size; - match.token_count = token_count; - return match; + 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; } -} // namespace - static void build_qwen_matmul_dispatch(const QwenMatmulMatch & match, DispatchMatch & dispatch_match, size_t root_index) { Dispatch dispatch; dispatch.kernel = make_kernel_specialization(match.kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - if (match.kernel.id == kGgmlLinearQ6KQ8_1X4Kernel.id) { - dispatch.kernel.integer_parameters.emplace("input_size", match.input_size); - dispatch.kernel.integer_parameters.emplace("output_size", match.output_size); - dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.token_capacity", - to_config_value(match.token_count)); - dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.output_capacity", - to_config_value(match.output_size)); - } else { - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", - to_config_value(match.token_count)); - } - if (match.dense) { - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", - to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", - to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "0"); - } + dispatch.kernel.integer_parameters.emplace("input_size", match.input_size); + dispatch.kernel.integer_parameters.emplace("output_size", match.output_size); + dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.token_capacity", + to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace("ggml.linear_q6k_q8_1_x4.output_capacity", + to_config_value(match.output_size)); dispatch.bindings.push_back({ match.input_value, 0, match.input_bytes }); dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); @@ -481,16 +365,6 @@ static bool match_qwen_q6k_q8_dispatch(const DispatchMatchContext & context, Dis return true; } -static bool match_qwen_decode_endpoint_q6k_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { - const QwenMatmulMatch match = match_qwen_decode_endpoint_q6k_matmul(context.graph, context.root_node); - if (!match.matched()) { - return false; - } - build_qwen_matmul_dispatch(match, dispatch_match, context.root_index); - return true; -} - static bool match_qwen_attention_output_next_q8_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const QwenAttentionOutputNextQ8Match match = match_qwen_attention_output_next_q8(context); @@ -498,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); @@ -519,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 && @@ -558,70 +438,32 @@ static bool match_qwen_attention_output_next_q8_dispatch(const DispatchMatchCont return true; } -static bool match_qwen_attention_output_accumulate_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { - const QwenAttentionOutputAccumulateMatch match = match_qwen_attention_output_accumulate(context); - if (!match.matched()) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kQwenDenseLinearQ4KF16WmmaKernel); - dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.input_size", - to_config_value(match.input_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_size", - to_config_value(match.output_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.dense_quantized.output_accumulation", "1"); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); - dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); - dispatch.bindings.push_back({ match.residual_output->id, 0, match.residual_output->byte_count }); - - dispatch_match.value_aliases.push_back({ match.residual_input->id, match.residual_output->id }); - if (!append_covered_node_index_once(context.graph, context.covered_nodes, context.root_node, - dispatch_match.covered_nodes) || - !append_covered_node_index_once(context.graph, context.covered_nodes, match.add_node, - dispatch_match.covered_nodes)) { - return false; - } - dispatch_match.dispatches.push_back(std::move(dispatch)); - return true; -} - void register_qwen_matmul_dispatches(DispatchRegistryBuilder & registry) { - registry.add({ - "qwen.matmul.attention_output_q4k_q8_1_x4_next_q8", + registry.add_activation_consumer({ + "qwen.matmul.input.q8_1_x4.q6k_vocab", GGML_OP_MUL_MAT, - DispatchMatchKind::Fused, - 300, + DispatchActivationInputQ8_1X4, + DispatchActivationConsumerUse::BandwidthLimited, + DispatchActivationProducerRequirementNone, DispatchSource::Qwen, - match_qwen_attention_output_next_q8_dispatch, + qwen_accepts_q6k_q8_input, }); registry.add({ - "qwen.matmul.attention_output_q4k_f16_accumulate", + "qwen.matmul.attention_output_q4k_q8_1_x4_next_q8", GGML_OP_MUL_MAT, DispatchMatchKind::Fused, - 250, + 300, DispatchSource::Qwen, - match_qwen_attention_output_accumulate_dispatch, + match_qwen_attention_output_next_q8_dispatch, }); registry.add({ "qwen.matmul.q6k_q8_1_x4", GGML_OP_MUL_MAT, - DispatchMatchKind::SingleOp, - 200, + DispatchMatchKind::Fused, + 305, DispatchSource::Qwen, match_qwen_q6k_q8_dispatch, }); - registry.add({ - "qwen.matmul.decode_endpoint_q6k_f16_wmma", - GGML_OP_MUL_MAT, - DispatchMatchKind::SingleOp, - 100, - DispatchSource::Qwen, - match_qwen_decode_endpoint_q6k_dispatch, - }); } } // namespace ggml::hrx 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 new file mode 100644 index 000000000000..28f348e9cacb --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-matmul.h @@ -0,0 +1,13 @@ +#pragma once + +#include "../dispatch-registry.h" + +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/dispatch-rmsnorm.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp similarity index 82% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.cpp rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp index a28032969d1f..dce10fa60c64 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.cpp @@ -1,18 +1,20 @@ -#include "dispatch-rmsnorm.h" +#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" #include #include +#include #include #include namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kQwenRmsNormF32Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_rmsnorm_f32"); static constexpr KernelCatalogRef kQwenRmsNormF32QuantizeQ8_1X4Kernel = GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); static constexpr KernelCatalogRef kGgmlLinearQ6KQ8_1X4Kernel = @@ -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; @@ -242,40 +232,15 @@ static std::string to_config_value(int64_t value) { return std::to_string(value); } -} // namespace - -static bool match_qwen_rmsnorm_f32_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { - const std::vector & nodes = context.graph.nodes(); - if (context.root_index >= nodes.size()) { - return false; - } - const RmsNormMatch rms_match = - match_qwen_rmsnorm_f32(context.graph, &nodes[context.root_index], context.root_index); - if (!rms_match.matched() || rms_match.rms_node_index >= context.covered_nodes.size() || - rms_match.mul_node_index >= context.covered_nodes.size() || context.covered_nodes[rms_match.rms_node_index] || - context.covered_nodes[rms_match.mul_node_index]) { - return false; - } - - Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kQwenRmsNormF32Kernel); - dispatch.kernel.integer_parameters.emplace("token_count", rms_match.token_count); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.hidden_size", to_config_value(rms_match.hidden_size)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.model.rms_epsilon", "0.000001"); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", - to_config_value(rms_match.token_count)); - dispatch.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", - to_config_value(rms_match.q8_group_count)); - 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 }); - - 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)); - return true; +static std::string to_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); } +} // namespace + static bool match_qwen_rmsnorm_f32_quantize_q8_1_x4_dispatch(const DispatchMatchContext & context, DispatchMatch & match) { const std::vector & nodes = context.graph.nodes(); @@ -355,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); @@ -378,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(); } @@ -420,14 +385,6 @@ void register_qwen_rmsnorm_dispatches(DispatchRegistryBuilder & registry) { DispatchSource::Qwen, match_qwen_rmsnorm_f32_quantize_q8_1_x4_dispatch, }); - registry.add({ - "qwen.rmsnorm_f32.mul_weight", - GGML_OP_RMS_NORM, - DispatchMatchKind::Fused, - 1000, - DispatchSource::Qwen, - match_qwen_rmsnorm_f32_dispatch, - }); } } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.h similarity index 80% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.h rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.h index 4a238a7e22b8..e728f90a63dc 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-rmsnorm.h +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen-rmsnorm.h @@ -1,6 +1,6 @@ #pragma once -#include "dispatch-registry.h" +#include "../dispatch-registry.h" namespace ggml::hrx { diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen.cpp new file mode 100644 index 000000000000..28dad1bfff41 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen.cpp @@ -0,0 +1,19 @@ +#include "dispatch-qwen.h" + +#include "dispatch-moe-router.h" +#include "dispatch-qwen-attention-postprocess.h" +#include "dispatch-qwen-matmul.h" +#include "dispatch-qwen-rmsnorm.h" +#include "dispatch-routed-ffn.h" + +namespace ggml::hrx { + +void register_qwen_dispatches(DispatchRegistryBuilder & registry) { + register_qwen_attention_postprocess_dispatches(registry); + register_qwen_matmul_dispatches(registry); + register_routed_ffn_dispatches(registry); + register_qwen_rmsnorm_dispatches(registry); + register_moe_router_dispatches(registry); +} + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen.h new file mode 100644 index 000000000000..377477b7f7a5 --- /dev/null +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-qwen.h @@ -0,0 +1,9 @@ +#pragma once + +#include "../dispatch-registry.h" + +namespace ggml::hrx { + +void register_qwen_dispatches(DispatchRegistryBuilder & registry); + +} // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.cpp b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp similarity index 68% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.cpp rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp index 46bf7025649c..ef1d99e35ec2 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.cpp @@ -1,8 +1,12 @@ #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" +#include "graph/op-params.h" #include "kernel-corpus/kernel-corpus-catalog-verify.h" #include @@ -14,35 +18,35 @@ namespace ggml::hrx { namespace { -static constexpr KernelCatalogRef kQwenRoutedGateUpSwiGLUQ4KF16WmmaKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma"); -static constexpr KernelCatalogRef kQwenRoutedGateUpSwiGLUQ4KQ8Kernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_gate_up_swiglu_q4k_q8"); -static constexpr KernelCatalogRef kQwenRoutedGateUpSwiGLUQ4KQ8NextQ8Kernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8"); -static constexpr KernelCatalogRef kQwenRoutedDownQ4KF16WmmaGroupedKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q4k_f16_wmma_grouped"); -static constexpr KernelCatalogRef kQwenRoutedDownQ6KF16WmmaGroupedKernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q6k_f16_wmma_grouped"); -static constexpr KernelCatalogRef kQwenRoutedDownQ4KQ8NextQ8Kernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q4k_q8_1_x4_next_q8"); -static constexpr KernelCatalogRef kQwenRoutedDownQ6KF32Wave64NextQ8Kernel = - GGML_HRX_KERNEL_REF("qwen3_moe", "qwen3_moe_routed_down_q6k_f32_wave64_next_q8"); +static constexpr KernelCatalogRef kCommonRoutedGateUpSwiGLUF16WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_swiglu_f16_f16_wmma"); +static constexpr KernelCatalogRef kCommonVectorRoutedGateUpSwiGLUPublishF32Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_vector_pair_binary_publish_f32"); +static constexpr KernelCatalogRef kCommonVectorRoutedGateUpSwiGLUPublishQ8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_vector_pair_binary_publish_q8"); +static constexpr KernelCatalogRef kCommonMulMatIdF16F16WmmaKernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_f16_f16_wmma"); +static constexpr KernelCatalogRef kCommonVectorRoutedDownQ4KPublishQ8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_vector_weighted_wave32_publish_q8"); +static constexpr KernelCatalogRef kCommonVectorRoutedDownQ6KPublishQ8Kernel = + GGML_HRX_KERNEL_REF("loom_libs", "ggml_mul_mat_id_vector_weighted_wave64_publish_q8"); 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,13 +65,66 @@ 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; } static bool is_routed_ffn_gate_up_weight(const Value & value) { - return value.type == GGML_TYPE_Q4_K && value.contiguous && + return (value.type == GGML_TYPE_Q4_K || value.type == GGML_TYPE_Q6_K) && value.contiguous && + is_shape(value, kRoutedFfnInputSize, kRoutedFfnExpertHiddenSize, kRoutedFfnExpertCount, 1); +} + +static bool is_fast_routed_ffn_gate_up_weight(const Value & value) { + const bool supported_type = value.type == GGML_TYPE_Q4_K || value.type == GGML_TYPE_Q6_K || + value.type == GGML_TYPE_Q8_0 || value.type == GGML_TYPE_Q8_1 || + value.type == GGML_TYPE_F16; + return supported_type && value.contiguous && is_shape(value, kRoutedFfnInputSize, kRoutedFfnExpertHiddenSize, kRoutedFfnExpertCount, 1); } @@ -115,6 +172,16 @@ static const GraphNode * producer_with_op(const Graph & graph, ValueId value, gg return producer != nullptr && producer->op == op ? producer : nullptr; } +static bool is_swiglu_params(const OpParams & params) { + const BinaryParams * binary_params = op_params_as(params); + if (binary_params != nullptr) { + return binary_params->op == BinaryKind::SwiGLU; + } + + const GluParams * glu_params = op_params_as(params); + return glu_params != nullptr && glu_params->op == GGML_GLU_OP_SWIGLU; +} + static bool append_covered_node(const DispatchMatchContext & context, const GraphNode * node, DispatchMatch & match) { return append_covered_node_index_once(context.graph, context.covered_nodes, node, match.covered_nodes); } @@ -165,6 +232,60 @@ static void add_routed_down_compile_parameters(Dispatch & dispatch, int64_t toke dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(token_count)); } +static void add_common_vector_shape_compile_parameters(Dispatch & dispatch, int64_t input_size, int64_t output_size) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.input_size", to_config_value(input_size)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.route_count", + to_config_value(kRoutedFfnRouteCount)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.expert_count", + to_config_value(kRoutedFfnExpertCount)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.output_size", to_config_value(output_size)); +} + +static void add_common_vector_pair_compile_parameters(Dispatch & dispatch, int64_t token_count, ggml_type publish_type) { + add_common_vector_shape_compile_parameters(dispatch, kRoutedFfnInputSize, kRoutedFfnExpertHiddenSize); + const std::string q4k_format = + to_config_value(common_mul_mat_format_config_value(CommonMulMatWeightFormat::Q4K)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.input_format", + to_config_value(static_cast(GGML_TYPE_Q8_1))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.lhs_weight_format", q4k_format); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.rhs_weight_format", q4k_format); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_id.vector.binary_op", + to_config_value(static_cast(binary_kind_config_value(BinaryKind::SwiGLU)))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.publish_format", + to_config_value(static_cast(publish_type))); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", to_config_value(token_count)); +} + +static void add_common_vector_weighted_compile_parameters(Dispatch & dispatch, int64_t token_count, + ggml_type input_type, ggml_type weight_type) { + add_common_vector_shape_compile_parameters(dispatch, kRoutedFfnExpertHiddenSize, kRoutedFfnInputSize); + CommonMulMatWeightFormat weight_format; + const bool has_weight_format = common_mul_mat_format_for_type(weight_type, weight_format); + GGML_ASSERT(has_weight_format); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.input_format", + to_config_value(static_cast(input_type))); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat_id.vector.weight_format", to_config_value(common_mul_mat_format_config_value(weight_format))); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id.vector.publish_format", + to_config_value(static_cast(GGML_TYPE_Q8_1))); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", to_config_value(token_count)); +} + +static void add_common_routed_down_compile_parameters(Dispatch & dispatch, int64_t token_count, ggml_type weight_type) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_f16_f16.input_size", + to_config_value(kRoutedFfnExpertHiddenSize)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_f16_f16.route_count", + to_config_value(kRoutedFfnRouteCount)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_f16_f16.expert_count", + to_config_value(kRoutedFfnExpertCount)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_f16_f16.output_size", + to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_id_f16_f16.weight_format", + weight_type == GGML_TYPE_Q4_K ? "4" : "6"); + dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", to_config_value(token_count)); +} + struct RoutedGateUpMatch { const Value * gate_weight = nullptr; const Value * up_weight = nullptr; @@ -305,6 +426,65 @@ 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, + ValueId expected_route_ids, + int64_t token_count, + const Value *& weight, + const Value *& output) { + if (node.op != GGML_OP_MUL_MAT_ID || node.inputs.size() != 3 || node.inputs[1] != expected_input || + node.inputs[2] != expected_route_ids) { + return false; + } + weight = graph_value(graph, node.inputs[0]); + output = graph_value(graph, node.output); + return weight != nullptr && output != nullptr && is_fast_routed_ffn_gate_up_weight(*weight) && + is_routed_ffn_projection_output(*output, token_count); +} + static RoutedDownMatch match_routed_ffn_down_grouped(const DispatchMatchContext & context) { RoutedDownMatch match; const GraphNode * root = context.root_node; @@ -334,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 {}; } @@ -345,8 +526,7 @@ static RoutedDownMatch match_routed_ffn_down_grouped(const DispatchMatchContext match.output = root_output; match.route_ids = route_ids; match.routing_bundle = routing_bundle; - match.kernel = weight->type == GGML_TYPE_Q4_K ? kQwenRoutedDownQ4KF16WmmaGroupedKernel : - kQwenRoutedDownQ6KF16WmmaGroupedKernel; + match.kernel = kCommonMulMatIdF16F16WmmaKernel; match.token_count = token_count; return match; } @@ -363,7 +543,7 @@ static RoutedGateUpMatch match_routed_ffn_gate_up_swiglu(const DispatchMatchCont const Value * route_ids = graph_value(context.graph, root->inputs[2]); const Value * root_output = graph_value(context.graph, root->output); if (root_weight == nullptr || input == nullptr || route_ids == nullptr || root_output == nullptr || - !is_routed_ffn_gate_up_weight(*root_weight) || input->type != GGML_TYPE_F32 || !input->contiguous || + !is_fast_routed_ffn_gate_up_weight(*root_weight) || input->type != GGML_TYPE_F32 || !input->contiguous || !is_shape(*input, kRoutedFfnInputSize, 1, input->ne[2], 1) || route_ids->type != GGML_TYPE_I32 || !is_shape(*route_ids, kRoutedFfnRouteCount, input->ne[2], 1, 1)) { return {}; @@ -384,8 +564,7 @@ static RoutedGateUpMatch match_routed_ffn_gate_up_swiglu(const DispatchMatchCont if (glu_node == nullptr || glu_node->inputs.size() != 2) { return {}; } - const GluParams * glu_params = op_params_as(glu_node->params); - if (glu_params == nullptr || glu_params->op != GGML_GLU_OP_SWIGLU) { + if (!is_swiglu_params(glu_node->params)) { return {}; } @@ -399,10 +578,10 @@ static RoutedGateUpMatch match_routed_ffn_gate_up_swiglu(const DispatchMatchCont const Value * gate_output = nullptr; const Value * up_weight = nullptr; const Value * up_output = nullptr; - if (!match_same_route_projection(context.graph, *gate_node, input->id, route_ids->id, token_count, gate_weight, - gate_output) || - !match_same_route_projection(context.graph, *up_node, input->id, route_ids->id, token_count, up_weight, - up_output)) { + if (!match_same_fast_route_projection(context.graph, *gate_node, input->id, route_ids->id, token_count, gate_weight, + gate_output) || + !match_same_fast_route_projection(context.graph, *up_node, input->id, route_ids->id, token_count, up_weight, + up_output)) { return {}; } if (!same_shape(*gate_output, *up_output)) { @@ -459,8 +638,7 @@ static DecodeRoutedGateUpMatch match_decode_routed_ffn_gate_up_swiglu(const Disp if (glu_node == nullptr || glu_node->inputs.size() != 2) { return {}; } - const GluParams * glu_params = op_params_as(glu_node->params); - if (glu_params == nullptr || glu_params->op != GGML_GLU_OP_SWIGLU) { + if (!is_swiglu_params(glu_node->params)) { return {}; } @@ -577,6 +755,31 @@ static bool node_is_covered(const DispatchMatchContext & context, const GraphNod context.covered_nodes[index]; } +static bool has_uncovered_q6_k_final_projection_consumer(const DispatchMatchContext & context, ValueId input) { + for (const GraphNode * consumer : consumers_with_op_through_layout_aliases(context.graph, input, GGML_OP_MUL_MAT)) { + if (consumer == nullptr || node_is_covered(context, consumer) || consumer->inputs.size() < 2 || + !node_has_input_or_alias(context.graph, *consumer, input)) { + continue; + } + + const Value * weight = graph_value(context.graph, consumer->inputs[0]); + const Value * rhs = graph_value(context.graph, consumer->inputs[1]); + const Value * output = graph_value(context.graph, consumer->output); + if (weight == nullptr || rhs == nullptr || output == nullptr || weight->type != GGML_TYPE_Q6_K || + rhs->type != GGML_TYPE_F32 || output->type != GGML_TYPE_F32 || weight->alias_source.value >= 0 || + rhs->ne[0] <= 0 || output->ne[0] <= 0) { + continue; + } + + const int64_t input_size = rhs->ne[0]; + const int64_t output_size = output->ne[0]; + if (input_size % 256 == 0 && output_size % 64 == 0 && output_size >= 16 * input_size) { + return true; + } + } + return false; +} + static bool residual_input_is_safe_for_in_place(const DispatchMatchContext & context, const WeightedReduceMatch & match) { if (match.residual_input == nullptr || match.residual == nullptr) { @@ -791,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 {}; } @@ -812,36 +1016,51 @@ static WeightedReduceMatch match_routed_ffn_down_weighted_reduce(const DispatchM return match; } -static bool match_routed_ffn_gate_up_swiglu_q4k_f16_wmma_dispatch(const DispatchMatchContext & context, - DispatchMatch & dispatch_match) { +static bool match_routed_ffn_gate_up_swiglu_f16_wmma_dispatch(const DispatchMatchContext & context, + DispatchMatch & dispatch_match) { const RoutedGateUpMatch match = match_routed_ffn_gate_up_swiglu(context); if (!match.matched()) { return false; } + if (match.gate_weight->type != match.up_weight->type) { + return false; + } - 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)) { + CommonMulMatWeightFormat weight_format; + if (!common_mul_mat_format_for_type(match.gate_weight->type, weight_format)) { + return false; + } + 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 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; } Dispatch dispatch; - dispatch.kernel = make_kernel_specialization(kQwenRoutedGateUpSwiGLUQ4KF16WmmaKernel); + dispatch.kernel = make_kernel_specialization(kCommonRoutedGateUpSwiGLUF16WmmaKernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.input_size", - to_config_value(kRoutedFfnInputSize)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.expert_count", - to_config_value(kRoutedFfnExpertCount)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.route_count", - to_config_value(kRoutedFfnRouteCount)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.output_size", + dispatch.kernel.compile_parameters.emplace(config_prefix + ".input_size", to_config_value(kRoutedFfnInputSize)); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".expert_count", to_config_value(kRoutedFfnExpertCount)); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".route_count", to_config_value(kRoutedFfnRouteCount)); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".output_size", to_config_value(kRoutedFfnExpertHiddenSize)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".gate_weight_format", weight_format_config); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".up_weight_format", weight_format_config); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".descriptor_expert_mask", "127"); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".descriptor_partition_shift", "7"); + dispatch.kernel.compile_parameters.emplace(config_prefix + ".descriptor_row_count_shift", "13"); dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); dispatch.bindings.push_back( @@ -850,7 +1069,7 @@ static bool match_routed_ffn_gate_up_swiglu_q4k_f16_wmma_dispatch(const Dispatch { 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) || @@ -868,28 +1087,35 @@ 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 ? kQwenRoutedGateUpSwiGLUQ4KQ8NextQ8Kernel : - kQwenRoutedGateUpSwiGLUQ4KQ8Kernel); + dispatch.kernel = make_kernel_specialization(match.publish_q8 ? kCommonVectorRoutedGateUpSwiGLUPublishQ8Kernel : + kCommonVectorRoutedGateUpSwiGLUPublishF32Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.integer_parameters.emplace("route_count", kRoutedFfnRouteCount); dispatch.kernel.integer_parameters.emplace("route_stride", match.route_stride); dispatch.kernel.integer_parameters.emplace("expert_count", kRoutedFfnExpertCount); dispatch.kernel.integer_parameters.emplace("output_size", kRoutedFfnExpertHiddenSize); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.input_size", - to_config_value(kRoutedFfnInputSize)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.expert_count", - to_config_value(kRoutedFfnExpertCount)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.route_count", - to_config_value(kRoutedFfnRouteCount)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.routed_gate_up.output_size", - to_config_value(kRoutedFfnExpertHiddenSize)); - dispatch.kernel.compile_parameters.emplace("qwen3_moe.workload.token_capacity", to_config_value(match.token_count)); + add_common_vector_pair_compile_parameters(dispatch, match.token_count, + match.publish_q8 ? GGML_TYPE_Q8_1 : GGML_TYPE_F32); const size_t route_id_length = static_cast(match.token_count * match.route_stride) * sizeof(int32_t); dispatch.bindings.push_back({ match.input_alternate->alternate_value, 0, match.input_alternate->byte_count }); @@ -900,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) || @@ -970,6 +1187,9 @@ static DecodeRoutedDownMatch match_decode_routed_ffn_down_next_q8(const Dispatch !residual_input_is_safe_for_in_place(context, reduce)) { return {}; } + if (has_uncovered_q6_k_final_projection_consumer(context, reduce.next_rmsnorm.output->id)) { + return {}; + } const bool input_is_q8 = weight->type == GGML_TYPE_Q4_K; const CommandPlanAlternateValue * input_alternate = @@ -986,8 +1206,8 @@ static DecodeRoutedDownMatch match_decode_routed_ffn_down_next_q8(const Dispatch match.output = reduce.output; match.route_ids = route_ids; match.reduce = std::move(reduce); - match.kernel = - weight->type == GGML_TYPE_Q4_K ? kQwenRoutedDownQ4KQ8NextQ8Kernel : kQwenRoutedDownQ6KF32Wave64NextQ8Kernel; + match.kernel = weight->type == GGML_TYPE_Q4_K ? kCommonVectorRoutedDownQ4KPublishQ8Kernel : + kCommonVectorRoutedDownQ6KPublishQ8Kernel; match.token_count = 1; match.route_stride = route_stride; match.input_is_q8 = input_is_q8; @@ -1001,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); @@ -1013,7 +1254,9 @@ static bool match_decode_routed_ffn_down_next_q8_dispatch(const DispatchMatchCon dispatch.kernel.integer_parameters.emplace("route_id_stride", match.route_stride); dispatch.kernel.integer_parameters.emplace("expert_count", kRoutedFfnExpertCount); dispatch.kernel.integer_parameters.emplace("output_size", kRoutedFfnInputSize); - add_routed_down_compile_parameters(dispatch, match.token_count); + add_common_vector_weighted_compile_parameters(dispatch, match.token_count, + match.input_is_q8 ? GGML_TYPE_Q8_1 : GGML_TYPE_F32, + match.weight->type); 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"); @@ -1030,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; @@ -1075,32 +1302,37 @@ static bool match_decode_routed_ffn_down_next_q8_dispatch(const DispatchMatchCon static bool build_routed_ffn_down_grouped_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match, - KernelCatalogRef expected_kernel) { + ggml_type expected_weight_type) { const RoutedDownMatch match = match_routed_ffn_down_grouped(context); - if (!match.matched() || match.kernel.id != expected_kernel.id) { + if (!match.matched() || match.weight->type != expected_weight_type) { 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; } Dispatch dispatch; dispatch.kernel = make_kernel_specialization(match.kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); - add_routed_down_compile_parameters(dispatch, match.token_count); + add_common_routed_down_compile_parameters(dispatch, match.token_count, match.weight->type); dispatch.bindings.push_back({ match.input_alternate->alternate_value, 0, match.input_alternate->byte_count }); 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)); @@ -1116,24 +1348,54 @@ 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 }); + dispatch.bindings.push_back({ match.residual_input->id, 0, match.residual_input->byte_count }); dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); } @@ -1164,17 +1426,53 @@ static bool match_routed_ffn_down_weighted_reduce_dispatch(const DispatchMatchCo static bool match_routed_ffn_down_q4k_f16_wmma_grouped_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - return build_routed_ffn_down_grouped_dispatch(context, dispatch_match, kQwenRoutedDownQ4KF16WmmaGroupedKernel); + return build_routed_ffn_down_grouped_dispatch(context, dispatch_match, GGML_TYPE_Q4_K); } static bool match_routed_ffn_down_q6k_f16_wmma_grouped_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - return build_routed_ffn_down_grouped_dispatch(context, dispatch_match, kQwenRoutedDownQ6KF16WmmaGroupedKernel); + return build_routed_ffn_down_grouped_dispatch(context, dispatch_match, GGML_TYPE_Q6_K); } } // 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, @@ -1192,12 +1490,12 @@ void register_routed_ffn_dispatches(DispatchRegistryBuilder & registry) { match_decode_routed_ffn_down_next_q8_dispatch, }); registry.add({ - "llm.routed_ffn.gate_up_swiglu_q4k_f16_wmma", + "llm.routed_ffn.gate_up_swiglu_f16_wmma", GGML_OP_MUL_MAT_ID, DispatchMatchKind::Fused, 1000, DispatchSource::Llm, - match_routed_ffn_gate_up_swiglu_q4k_f16_wmma_dispatch, + match_routed_ffn_gate_up_swiglu_f16_wmma_dispatch, }); registry.add({ "llm.routed_ffn.down_q4k_f16_wmma_grouped", diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.h similarity index 80% rename from ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.h rename to ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.h index 18cdabc5b51f..54589e193b73 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-routed-ffn.h +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-routed-ffn.h @@ -1,6 +1,6 @@ #pragma once -#include "dispatch-registry.h" +#include "../dispatch-registry.h" namespace ggml::hrx { diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 8b1fe8bc669b..81ecb97e514f 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -4,6 +4,7 @@ #include "backend-context.h" #include "ggml-backend-impl.h" #include "ggml-impl.h" +#include "graph/op-params.h" #include "hrx_runtime.h" #include "kernel-corpus/kernel-corpus.h" #include "loom-jit.h" @@ -19,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -40,6 +42,22 @@ static bool environment_flag_enabled(const char * name) { return value != nullptr && value[0] != '\0' && value[0] != '0'; } +static void log_hrx_device_event(void *, const hrx_device_event_t * event) { + if (event == nullptr || event->type != HRX_DEVICE_EVENT_TYPE_ASAN_REPORT || event->payload.data == nullptr || + event->payload.data_length < sizeof(hrx_device_asan_report_t)) { + return; + } + hrx_device_asan_report_t report; + std::memcpy(&report, event->payload.data, sizeof(report)); + GGML_LOG_ERROR("HRX ASAN: executable=%" PRIu64 " export=%u site=%" PRIu64 " access=%u address=0x%016" PRIx64 + " length=%" PRIu64 " workgroup=(%u,%u,%u) workitem=(%u,%u,%u) shadow=0x%016" PRIx64 + " value=0x%016" PRIx64 " dispatch=0x%016" PRIx64 "\n", + event->source.executable_id, event->source.export_ordinal, report.site_id, report.access_kind, + report.fault_address, report.access_length, report.workgroup_id[0], report.workgroup_id[1], + report.workgroup_id[2], report.workitem_id[0], report.workitem_id[1], report.workitem_id[2], + report.shadow_address, report.shadow_value, report.source_dispatch_ptr); +} + static bool hrx_check(hrx_status_t status, const char * expression, const char * file, int line) { if (hrx_status_is_ok(status)) { return true; @@ -278,10 +296,10 @@ static const ggml_backend_buffer_i buffer_i = { }; static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_t size) { - auto * type_context = static_cast(buft->context); - const bool host_visible = type_context->host_visible; - const bool direct_host_binding = host_visible && type_context->device->use_direct_host_bindings; - hrx_memory_type_t memory_type = HRX_MEMORY_TYPE_DEVICE_LOCAL; + auto * type_context = static_cast(buft->context); + const bool host_visible = type_context->host_visible; + const bool direct_host_binding = host_visible && type_context->device->use_direct_host_bindings; + hrx_memory_type_t memory_type = HRX_MEMORY_TYPE_DEVICE_LOCAL; // Direct command-program bindings require coherent CPU/GPU visibility. Otherwise HRX host buffers are pinned // transfer memory: DEVICE_VISIBLE permits handle-based stream copies without implying direct device access. if (host_visible) { @@ -290,7 +308,7 @@ static ggml_backend_buffer_t buffer_alloc(ggml_backend_buffer_type_t buft, size_ memory_type |= HRX_MEMORY_TYPE_HOST_COHERENT; } } - hrx_buffer_params_t params = { + hrx_buffer_params_t params = { memory_type, HRX_MEMORY_ACCESS_ALL, host_visible ? @@ -345,9 +363,11 @@ static const char * backend_name(ggml_backend_t backend) { return static_cast(backend->context)->name.c_str(); } +static void backend_synchronize(ggml_backend_t backend); + static void backend_free(ggml_backend_t backend) { auto * context = static_cast(backend->context); - HRX_CHECK(hrx_stream_synchronize(context->stream)); + backend_synchronize(backend); context->prepared_programs.clear(); context->graph_programs.clear(); context->kernel_executables.clear(); @@ -364,11 +384,11 @@ static bool synchronous_upload_fallback(ggml_backend_hrx_context * backend, hrx_buffer_t destination, size_t destination_offset, size_t size) { - const uint64_t fallback = - backend->device->synchronous_upload_fallbacks.fetch_add(1, std::memory_order_relaxed); + const uint64_t fallback = backend->device->synchronous_upload_fallbacks.fetch_add(1, std::memory_order_relaxed); if (fallback == 0) { - GGML_LOG_WARN("ggml_hrx: synchronous upload fallback for an unregistered host pointer; use the HRX host " - "buffer type for asynchronous transfers\n"); + GGML_LOG_WARN( + "ggml_hrx: synchronous upload fallback for an unregistered host pointer; use the HRX host " + "buffer type for asynchronous transfers\n"); } // Compatibility path for arbitrary GGML pointers. Keep the synchronization explicit until a bounded staging ring // with transfer retirement is available. @@ -386,11 +406,11 @@ static bool synchronous_download_fallback(ggml_backend_hrx_context * backend, size_t source_offset, void * destination, size_t size) { - const uint64_t fallback = - backend->device->synchronous_download_fallbacks.fetch_add(1, std::memory_order_relaxed); + const uint64_t fallback = backend->device->synchronous_download_fallbacks.fetch_add(1, std::memory_order_relaxed); if (fallback == 0) { - GGML_LOG_WARN("ggml_hrx: synchronous download fallback for an unregistered host pointer; use the HRX host " - "buffer type for asynchronous transfers\n"); + GGML_LOG_WARN( + "ggml_hrx: synchronous download fallback for an unregistered host pointer; use the HRX host " + "buffer type for asynchronous transfers\n"); } // Compatibility path for arbitrary GGML pointers. Keep the synchronization explicit until a bounded staging ring // with transfer retirement is available. @@ -472,15 +492,17 @@ static bool backend_copy_tensor_async(ggml_backend_t backend_src, } ggml_backend_buffer_t source_buffer = source->view_src != nullptr ? source->view_src->buffer : source->buffer; if (source_buffer != nullptr && ggml_backend_buffer_is_host(source_buffer)) { - return synchronous_upload_fallback( - destination_backend, source->data, destination_context->buffer, destination_offset, size); + return synchronous_upload_fallback(destination_backend, source->data, destination_context->buffer, + destination_offset, size); } return false; } static void backend_synchronize(ggml_backend_t backend) { auto * context = static_cast(backend->context); - HRX_CHECK(hrx_stream_synchronize(context->stream)); + if (HRX_CHECK(hrx_stream_synchronize(context->stream))) { + context->graph_replay_state.mark_stream_synchronized(); + } } static const char * status_first_error(const ggml::hrx::Status & status) { @@ -584,10 +606,15 @@ static bool eager_capability_declared(enum ggml_op op) { case GGML_OP_ADD: case GGML_OP_ARGSORT: case GGML_OP_CLAMP: + case GGML_OP_CONCAT: + case GGML_OP_CONT: + case GGML_OP_CPY: case GGML_OP_DIV: case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_GATED_DELTA_NET: case GGML_OP_GET_ROWS: case GGML_OP_GLU: + case GGML_OP_L2_NORM: case GGML_OP_MUL: case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: @@ -595,9 +622,13 @@ static bool eager_capability_declared(enum ggml_op op) { case GGML_OP_RESHAPE: case GGML_OP_RMS_NORM: case GGML_OP_ROPE: + case GGML_OP_SCALE: case GGML_OP_SET_ROWS: case GGML_OP_SOFT_MAX: + case GGML_OP_SSM_CONV: case GGML_OP_SUM_ROWS: + case GGML_OP_TRANSPOSE: + case GGML_OP_UNARY: case GGML_OP_VIEW: return true; default: @@ -605,9 +636,427 @@ static bool eager_capability_declared(enum ggml_op op) { } } +static bool zero_output_elision_safe_op(enum ggml_op op) { + switch (op) { + case GGML_OP_DUP: + case GGML_OP_ADD: + case GGML_OP_ADD_ID: + case GGML_OP_ADD1: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_SUM: + case GGML_OP_SUM_ROWS: + case GGML_OP_CUMSUM: + case GGML_OP_MEAN: + case GGML_OP_ARGMAX: + case GGML_OP_COUNT_EQUAL: + case GGML_OP_REPEAT: + case GGML_OP_REPEAT_BACK: + case GGML_OP_CONCAT: + case GGML_OP_SILU_BACK: + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_RMS_NORM_BACK: + case GGML_OP_GROUP_NORM: + case GGML_OP_L2_NORM: + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_OUT_PROD: + case GGML_OP_SCALE: + case GGML_OP_CONT: + case GGML_OP_GET_ROWS: + case GGML_OP_GET_ROWS_BACK: + case GGML_OP_DIAG: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_DIAG_MASK_ZERO: + case GGML_OP_SOFT_MAX: + case GGML_OP_SOFT_MAX_BACK: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + case GGML_OP_CLAMP: + case GGML_OP_CONV_TRANSPOSE_1D: + case GGML_OP_IM2COL: + case GGML_OP_IM2COL_BACK: + case GGML_OP_IM2COL_3D: + case GGML_OP_COL2IM_1D: + case GGML_OP_CONV_2D: + case GGML_OP_CONV_3D: + case GGML_OP_CONV_2D_DW: + case GGML_OP_CONV_TRANSPOSE_2D: + case GGML_OP_POOL_1D: + case GGML_OP_POOL_2D: + case GGML_OP_POOL_2D_BACK: + case GGML_OP_UPSCALE: + case GGML_OP_PAD: + case GGML_OP_PAD_REFLECT_1D: + case GGML_OP_ROLL: + case GGML_OP_ARANGE: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_ARGSORT: + case GGML_OP_TOP_K: + case GGML_OP_LEAKY_RELU: + case GGML_OP_TRI: + case GGML_OP_FILL: + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_FLASH_ATTN_BACK: + case GGML_OP_SSM_CONV: + case GGML_OP_SSM_SCAN: + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_GET_REL_POS: + case GGML_OP_ADD_REL_POS: + case GGML_OP_RWKV_WKV6: + case GGML_OP_GATED_LINEAR_ATTN: + case GGML_OP_RWKV_WKV7: + case GGML_OP_SOLVE_TRI: + case GGML_OP_GATED_DELTA_NET: + case GGML_OP_LIGHTNING_INDEXER: + case GGML_OP_DSV4_HC_COMB: + case GGML_OP_DSV4_HC_PRE: + case GGML_OP_DSV4_HC_POST: + case GGML_OP_UNARY: + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_GLU: + return true; + default: + return false; + } +} + +static bool zero_output_elision_supported(const ggml_tensor * op) { + return op != nullptr && (ggml_nelements(op) == 0 || ggml_nbytes(op) == 0) && zero_output_elision_safe_op(op->op); +} + +struct TensorStorageRange { + const ggml_tensor * root = nullptr; + size_t offset = 0; + size_t size = 0; + bool valid = false; +}; + +static TensorStorageRange tensor_storage_range(const ggml_tensor * tensor) { + TensorStorageRange range; + if (tensor == nullptr) { + return range; + } + range.size = ggml_nbytes(tensor); + while (tensor->view_src != nullptr) { + if (range.offset > std::numeric_limits::max() - tensor->view_offs) { + return {}; + } + range.offset += tensor->view_offs; + tensor = tensor->view_src; + } + range.root = tensor; + range.valid = true; + return range; +} + +static TensorStorageRange tensor_storage_range(const ggml_tensor * tensor, size_t size) { + TensorStorageRange range = tensor_storage_range(tensor); + if (range.valid) { + range.size = size; + } + return range; +} + +static bool tensor_storage_exactly_overlaps(const TensorStorageRange & lhs, const TensorStorageRange & rhs) { + return lhs.root == rhs.root && lhs.offset == rhs.offset && lhs.size == rhs.size; +} + +static bool tensor_storage_ranges_disjoint(const TensorStorageRange & lhs, const TensorStorageRange & rhs) { + if (lhs.root != rhs.root) { + return true; + } + if (lhs.offset > std::numeric_limits::max() - lhs.size || + rhs.offset > std::numeric_limits::max() - rhs.size) { + return false; + } + return lhs.offset + lhs.size <= rhs.offset || rhs.offset + rhs.size <= lhs.offset; +} + +static bool scale_storage_is_safe(const ggml_tensor * input, const ggml_tensor * output) { + const TensorStorageRange input_range = tensor_storage_range(input); + const TensorStorageRange output_range = tensor_storage_range(output); + if (!input_range.valid || !output_range.valid) { + return false; + } + return tensor_storage_exactly_overlaps(input_range, output_range) || + tensor_storage_ranges_disjoint(input_range, output_range); +} + +static bool binary_output_storage_is_safe(const ggml_tensor * lhs, const ggml_tensor * rhs, const ggml_tensor * output) { + const TensorStorageRange lhs_range = tensor_storage_range(lhs); + const TensorStorageRange rhs_range = tensor_storage_range(rhs); + const TensorStorageRange output_range = tensor_storage_range(output); + if (!lhs_range.valid || !rhs_range.valid || !output_range.valid) { + return false; + } + return tensor_storage_ranges_disjoint(lhs_range, output_range) && + tensor_storage_ranges_disjoint(rhs_range, output_range); +} + +static bool binary_output_storage_is_safe(const ggml_tensor * lhs, + size_t lhs_byte_count, + const ggml_tensor * rhs, + size_t rhs_byte_count, + const ggml_tensor * output) { + const TensorStorageRange lhs_range = tensor_storage_range(lhs, lhs_byte_count); + const TensorStorageRange rhs_range = tensor_storage_range(rhs, rhs_byte_count); + const TensorStorageRange output_range = tensor_storage_range(output); + if (!lhs_range.valid || !rhs_range.valid || !output_range.valid) { + return false; + } + return tensor_storage_ranges_disjoint(lhs_range, output_range) && + tensor_storage_ranges_disjoint(rhs_range, output_range); +} + +static bool tensor_has_positive_shape(const ggml_tensor * tensor) { + if (tensor == nullptr || ggml_nelements(tensor) <= 0) { + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] <= 0) { + return false; + } + } + return true; +} + +static bool tensor_has_packed_f32_layout(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return false; + } + size_t expected_stride = sizeof(float); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->nb[i] != expected_stride) { + return false; + } + expected_stride *= static_cast(tensor->ne[i]); + } + return true; +} + +static bool tensor_has_strided_f32_layout(const ggml_tensor * tensor) { + if (tensor == nullptr || tensor->nb[0] != sizeof(float)) { + return false; + } + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (tensor->nb[i] % sizeof(float) != 0) { + return false; + } + } + return true; +} + +static bool tensor_storage_span_bytes(const ggml_tensor * tensor, size_t & byte_count) { + if (!tensor_has_strided_f32_layout(tensor)) { + return false; + } + size_t max_offset = 0; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] <= 0) { + return false; + } + const size_t extent = static_cast(tensor->ne[i] - 1); + if (extent != 0 && tensor->nb[i] > std::numeric_limits::max() / extent) { + return false; + } + const size_t dim_offset = extent * tensor->nb[i]; + if (max_offset > std::numeric_limits::max() - dim_offset) { + return false; + } + max_offset += dim_offset; + } + if (max_offset > std::numeric_limits::max() - sizeof(float)) { + return false; + } + byte_count = max_offset + sizeof(float); + return true; +} + +static bool tensor_broadcastable_to(const ggml_tensor * source, const ggml_tensor * output) { + if (source == nullptr || output == nullptr) { + return false; + } + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (source->ne[i] != output->ne[i] && source->ne[i] != 1) { + return false; + } + } + return true; +} + +static bool binary_kind_allows_broadcast(ggml::hrx::BinaryKind kind, + const ggml_tensor * lhs, + const ggml_tensor * rhs, + const ggml_tensor * output) { + const bool lhs_full = ggml_are_same_shape(lhs, output); + const bool rhs_full = ggml_are_same_shape(rhs, output); + if (!tensor_broadcastable_to(lhs, output) || !tensor_broadcastable_to(rhs, output) || (!lhs_full && !rhs_full)) { + return false; + } + + switch (kind) { + case ggml::hrx::BinaryKind::Add: + case ggml::hrx::BinaryKind::Mul: + return true; + case ggml::hrx::BinaryKind::Sub: + case ggml::hrx::BinaryKind::Div: + return lhs_full; + case ggml::hrx::BinaryKind::SwiGLU: + case ggml::hrx::BinaryKind::GeGLU: + case ggml::hrx::BinaryKind::RegLU: + case ggml::hrx::BinaryKind::GeGLUErf: + case ggml::hrx::BinaryKind::GeGLUQuick: + return lhs_full && rhs_full; + } + return false; +} + +static bool supported_binary_f32_tensor(const ggml_tensor * op) { + if (op == nullptr || op->src[0] == nullptr || op->src[1] == nullptr) { + return false; + } + if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) { + return false; + } + if (!tensor_has_positive_shape(op)) { + return false; + } + if (!ggml_is_contiguous(op) || !tensor_has_packed_f32_layout(op) || op->view_src != nullptr) { + return false; + } + size_t lhs_byte_count = 0; + size_t rhs_byte_count = 0; + if (!tensor_storage_span_bytes(op->src[0], lhs_byte_count) || + !tensor_storage_span_bytes(op->src[1], rhs_byte_count)) { + return false; + } + if (!binary_output_storage_is_safe(op->src[0], lhs_byte_count, op->src[1], rhs_byte_count, op)) { + return false; + } + + ggml::hrx::BinaryKind binary_kind; + if (!ggml::hrx::import_binary_kind(*op, binary_kind)) { + return false; + } + + if (!ggml::hrx::binary_kind_supported(binary_kind)) { + return false; + } + + if ((!ggml_are_same_shape(op->src[0], op) || !ggml_are_same_shape(op->src[1], op)) && + (!ggml_is_contiguous(op->src[0]) || !ggml_is_contiguous(op->src[1]) || + !tensor_has_packed_f32_layout(op->src[0]) || !tensor_has_packed_f32_layout(op->src[1]))) { + return false; + } + + if (!binary_kind_allows_broadcast(binary_kind, op->src[0], op->src[1], op)) { + return false; + } + + return true; +} + +static bool supported_scale_f32_tensor(const ggml_tensor * op) { + return op != nullptr && op->op == GGML_OP_SCALE && op->src[0] != nullptr && op->type == GGML_TYPE_F32 && + op->src[0]->type == GGML_TYPE_F32 && ggml_are_same_shape(op, op->src[0]) && tensor_has_positive_shape(op) && + ggml_is_contiguous(op) && ggml_is_contiguous(op->src[0]) && tensor_has_packed_f32_layout(op) && + tensor_has_packed_f32_layout(op->src[0]) && scale_storage_is_safe(op->src[0], op) && + static_cast(ggml_nelements(op)) <= std::numeric_limits::max(); +} + +static bool supported_qwen_attention_projection_get_rows_tensor(const ggml_tensor * op) { + if (op == nullptr || op->op != GGML_OP_GET_ROWS || op->src[0] == nullptr || op->src[1] == nullptr || + op->src[0]->op != GGML_OP_MUL_MAT || op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || + op->src[1]->type != GGML_TYPE_I32 || !ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0]) || + !ggml_is_contiguous(op->src[1]) || op->ne[0] != op->src[0]->ne[0] || op->ne[1] <= 0 || + op->src[1]->ne[0] != op->ne[1] || op->ne[2] != 1 || op->ne[3] != 1) { + return false; + } + return true; +} + +static bool supported_qwen_attention_residual_add_tensor(const ggml_tensor * op) { + if (op == nullptr || op->op != GGML_OP_ADD || op->src[0] == nullptr || op->src[1] == nullptr || + op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 || + !ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0]) || !ggml_is_contiguous(op->src[1]) || + !ggml_are_same_shape(op, op->src[0]) || !ggml_are_same_shape(op, op->src[1])) { + return false; + } + + return op->src[0]->op == GGML_OP_GET_ROWS || op->src[1]->op == GGML_OP_GET_ROWS || + op->src[0]->op == GGML_OP_MUL_MAT || op->src[1]->op == GGML_OP_MUL_MAT; +} + +static bool supported_qwen_routed_ffn_reduce_add_tensor(const ggml_tensor * op) { + if (op == nullptr || op->op != GGML_OP_ADD || op->src[0] == nullptr || op->src[1] == nullptr || + op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 || + !ggml_is_contiguous(op) || !ggml_are_same_shape(op, op->src[0]) || !ggml_are_same_shape(op, op->src[1]) || + op->ne[0] != 2048 || op->ne[3] != 1 || + !((op->ne[1] > 0 && op->ne[2] == 1) || (op->ne[1] == 1 && op->ne[2] > 0))) { + return false; + } + + return op->src[0]->op == GGML_OP_VIEW || op->src[1]->op == GGML_OP_VIEW || op->src[0]->op == GGML_OP_ADD || + op->src[1]->op == GGML_OP_ADD; +} + +static bool supported_unary_f32_tensor(const ggml_tensor * op) { + if (op == nullptr || op->src[0] == nullptr || op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || + !ggml_are_same_shape(op, op->src[0]) || !ggml_is_contiguous(op) || !ggml_is_contiguous(op->src[0])) { + return false; + } + + ggml::hrx::UnaryKind unary_kind; + if (!ggml::hrx::import_unary_kind(*op, unary_kind)) { + return false; + } + + return ggml::hrx::unary_kind_supported(unary_kind); +} + static bool device_supports_op(ggml_backend_dev_t device, const ggml_tensor * op) { GGML_UNUSED(device); - return op != nullptr && eager_capability_declared(op->op); + if (op == nullptr) { + return false; + } + if (zero_output_elision_supported(op)) { + return true; + } + const bool supported_binary = supported_binary_f32_tensor(op); + if (supported_binary) { + return true; + } + if (supported_scale_f32_tensor(op)) { + return true; + } + if (supported_qwen_attention_projection_get_rows_tensor(op)) { + return true; + } + if (supported_qwen_attention_residual_add_tensor(op)) { + return true; + } + if (supported_qwen_routed_ffn_reduce_add_tensor(op)) { + return true; + } + const bool supported_unary = supported_unary_f32_tensor(op); + if (supported_unary) { + return true; + } + if (op->op == GGML_OP_GET_ROWS && op->src[0] != nullptr && + (op->src[0]->type == GGML_TYPE_IQ3_S || op->src[0]->type == GGML_TYPE_IQ4_NL)) { + return false; + } + return eager_capability_declared(op->op); } static bool device_supports_buffer_type(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { @@ -658,6 +1107,12 @@ static const ggml_backend_reg_i registry_i = { registry_name, registry_device_co static std::unique_ptr create_registry_context() { auto context = std::make_unique(); + if (environment_flag_enabled("GGML_HRX_LOG_DEVICE_EVENTS")) { + hrx_device_event_sink_t sink = { log_hrx_device_event, nullptr }; + if (!HRX_CHECK(hrx_runtime_set_device_event_sink(sink))) { + return context; + } + } hrx_status_t status = hrx_gpu_initialize(0); if (hrx_status_is_ok(status)) { context->initialized = true; diff --git a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp index 23292feef5e0..36df4bf405b8 100644 --- a/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp +++ b/ggml/src/ggml-hrx/graph/graph-diagnostics.cpp @@ -139,7 +139,24 @@ json op_params_json(const OpParams & params) { }; } else if constexpr (std::is_same_v) { return { - { "kind", "glu" }, + { "kind", "glu" }, + { "op", static_cast(value.op) }, + { "swapped", value.swapped } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "scale" }, + { "scale", value.scale }, + { "bias", value.bias } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "binary" }, + { "op", static_cast(value.op) } + }; + } else if constexpr (std::is_same_v) { + return { + { "kind", "unary" }, { "op", static_cast(value.op) } }; } else if constexpr (std::is_same_v) { @@ -154,6 +171,7 @@ json op_params_json(const OpParams & params) { { "attn_factor", value.attn_factor }, { "beta_fast", value.beta_fast }, { "beta_slow", value.beta_slow }, + { "sections", value.sections }, }; } }, @@ -187,13 +205,29 @@ OpParams parse_op_params(const json & item) { return ClampParams{ min, max }; } if (kind == "glu") { - return GluParams{ static_cast(item.value("op", static_cast(GGML_GLU_OP_REGLU))) }; + return GluParams{ static_cast(item.value("op", static_cast(GGML_GLU_OP_REGLU))), + item.value("swapped", false) }; + } + if (kind == "scale") { + return ScaleParams{ item.value("scale", 0.0f), item.value("bias", 0.0f) }; + } + if (kind == "binary") { + return BinaryParams{ static_cast(item.value("op", static_cast(BinaryKind::Add))) }; + } + if (kind == "unary") { + return UnaryParams{ static_cast(item.value("op", static_cast(UnaryKind::Abs))) }; } if (kind == "rope") { + std::array sections = {}; + if (item.contains("sections")) { + sections = item["sections"].get>(); + } return RopeParams{ - item.value("n_dims", 0), item.value("mode", 0), item.value("n_ctx_orig", 0), - item.value("freq_base", 0.0f), item.value("freq_scale", 0.0f), item.value("ext_factor", 0.0f), - item.value("attn_factor", 0.0f), item.value("beta_fast", 0.0f), item.value("beta_slow", 0.0f), + item.value("n_dims", 0), item.value("mode", 0), + item.value("n_ctx_orig", 0), item.value("freq_base", 0.0f), + item.value("freq_scale", 0.0f), item.value("ext_factor", 0.0f), + item.value("attn_factor", 0.0f), item.value("beta_fast", 0.0f), + item.value("beta_slow", 0.0f), sections, }; } return std::monostate{}; diff --git a/ggml/src/ggml-hrx/graph/graph.cpp b/ggml/src/ggml-hrx/graph/graph.cpp index 63221ff78dad..f88f34faba02 100644 --- a/ggml/src/ggml-hrx/graph/graph.cpp +++ b/ggml/src/ggml-hrx/graph/graph.cpp @@ -4,18 +4,48 @@ #include #include +#include #include namespace ggml::hrx { namespace { +static const ggml_tensor * tensor_storage_root(const ggml_tensor * tensor) { + while (tensor != nullptr && tensor->view_src != nullptr) { + tensor = tensor->view_src; + } + return tensor; +} + +static int32_t graph_tensor_use_count(const ggml_cgraph & graph, const ggml_tensor * tensor) { + if (tensor == nullptr || graph.use_counts == nullptr || graph.visited_hash_set.keys == nullptr) { + return -1; + } + const size_t hash_pos = ggml_hash_find(&graph.visited_hash_set, tensor); + if (hash_pos == GGML_HASHSET_FULL || !ggml_bitset_get(graph.visited_hash_set.used, hash_pos)) { + return -1; + } + return graph.use_counts[hash_pos]; +} + static bool tensor_is_external(const ggml_tensor * tensor, - const std::unordered_map & use_counts) { - if (tensor->op == GGML_OP_NONE) { + const std::unordered_map & use_counts, + const std::unordered_set & graph_nodes, + const ggml_cgraph & graph) { + const ggml_tensor * root = tensor_storage_root(tensor); + if (root == nullptr || root->op == GGML_OP_NONE || graph_nodes.find(root) == graph_nodes.end()) { + return true; + } + if (tensor->flags & GGML_TENSOR_FLAG_OUTPUT) { + return true; + } + const auto found = use_counts.find(tensor); + const int local_uses = found != use_counts.end() ? found->second : 0; + const int32_t graph_uses = graph_tensor_use_count(graph, tensor); + if (graph_uses > local_uses) { return true; } - const auto found = use_counts.find(tensor); - return found == use_counts.end() || found->second == 0; + return local_uses == 0; } } // namespace @@ -119,12 +149,15 @@ const GraphIndex & Graph::index() const { GraphImportResult import_ggml_graph(const ggml_cgraph & graph) { GraphImportResult result; std::unordered_map use_counts; + std::unordered_set graph_nodes; + graph_nodes.reserve(static_cast(graph.n_nodes)); for (int i = 0; i < graph.n_nodes; ++i) { const ggml_tensor * node = graph.nodes[i]; if (node == nullptr) { result.status.log("ggml graph contains a null node"); return result; } + graph_nodes.insert(node); for (const ggml_tensor * source : node->src) { if (source != nullptr) { ++use_counts[source]; @@ -140,14 +173,16 @@ GraphImportResult import_ggml_graph(const ggml_cgraph & graph) { if (source == nullptr) { continue; } - const ValueKind kind = tensor_is_external(source, use_counts) ? ValueKind::External : ValueKind::Transient; + const ValueKind kind = + tensor_is_external(source, use_counts, graph_nodes, graph) ? ValueKind::External : ValueKind::Transient; inputs.push_back(values.get_or_add_tensor_value(source, kind)); } - const ValueKind output_kind = tensor_is_external(node, use_counts) ? ValueKind::External : ValueKind::Transient; - const ValueId output = values.get_or_add_tensor_value(node, output_kind); - GraphNode & graph_node = result.graph.add_node(node->op, output, std::move(inputs)); - graph_node.params = import_op_params(*node); + const ValueKind output_kind = + tensor_is_external(node, use_counts, graph_nodes, graph) ? ValueKind::External : ValueKind::Transient; + const ValueId output = values.get_or_add_tensor_value(node, output_kind); + GraphNode & graph_node = result.graph.add_node(node->op, output, std::move(inputs)); + graph_node.params = import_op_params(*node); } result.status.append(result.graph.build_index()); diff --git a/ggml/src/ggml-hrx/graph/op-params.cpp b/ggml/src/ggml-hrx/graph/op-params.cpp index f76097164e14..015b6b7b77e3 100644 --- a/ggml/src/ggml-hrx/graph/op-params.cpp +++ b/ggml/src/ggml-hrx/graph/op-params.cpp @@ -51,6 +51,26 @@ static bool clamp_params_equivalent(const OpParams & lhs, const OpParams & rhs) static bool glu_params_equivalent(const OpParams & lhs, const OpParams & rhs) { const GluParams * lhs_params = op_params_as(lhs); const GluParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && lhs_params->op == rhs_params->op && + lhs_params->swapped == rhs_params->swapped; +} + +static bool scale_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const ScaleParams * lhs_params = op_params_as(lhs); + const ScaleParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && nearly_equal(lhs_params->scale, rhs_params->scale) && + nearly_equal(lhs_params->bias, rhs_params->bias); +} + +static bool binary_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const BinaryParams * lhs_params = op_params_as(lhs); + const BinaryParams * rhs_params = op_params_as(rhs); + return lhs_params != nullptr && rhs_params != nullptr && lhs_params->op == rhs_params->op; +} + +static bool unary_params_equivalent(const OpParams & lhs, const OpParams & rhs) { + const UnaryParams * lhs_params = op_params_as(lhs); + const UnaryParams * rhs_params = op_params_as(rhs); return lhs_params != nullptr && rhs_params != nullptr && lhs_params->op == rhs_params->op; } @@ -64,14 +84,180 @@ static bool rope_params_equivalent(const OpParams & lhs, const OpParams & rhs) { nearly_equal(lhs_params->ext_factor, rhs_params->ext_factor) && nearly_equal(lhs_params->attn_factor, rhs_params->attn_factor) && nearly_equal(lhs_params->beta_fast, rhs_params->beta_fast) && - nearly_equal(lhs_params->beta_slow, rhs_params->beta_slow); + nearly_equal(lhs_params->beta_slow, rhs_params->beta_slow) && lhs_params->sections == rhs_params->sections; } } // namespace +bool import_binary_kind(const ggml_tensor & tensor, BinaryKind & kind) { + switch (tensor.op) { + case GGML_OP_ADD: + kind = BinaryKind::Add; + return true; + case GGML_OP_SUB: + kind = BinaryKind::Sub; + return true; + case GGML_OP_MUL: + kind = BinaryKind::Mul; + return true; + case GGML_OP_DIV: + kind = BinaryKind::Div; + return true; + case GGML_OP_GLU: + if (tensor.src[1] == nullptr) { + return false; + } + switch (ggml_get_glu_op(&tensor)) { + case GGML_GLU_OP_REGLU: + kind = BinaryKind::RegLU; + return true; + case GGML_GLU_OP_SWIGLU: + kind = BinaryKind::SwiGLU; + return true; + case GGML_GLU_OP_GEGLU: + kind = BinaryKind::GeGLU; + return true; + case GGML_GLU_OP_GEGLU_ERF: + kind = BinaryKind::GeGLUErf; + return true; + case GGML_GLU_OP_GEGLU_QUICK: + kind = BinaryKind::GeGLUQuick; + return true; + default: + return false; + } + default: + return false; + } +} + +bool binary_kind_supported(BinaryKind kind) { + return static_cast(kind) <= static_cast(BinaryKind::GeGLUQuick); +} + +uint32_t binary_kind_config_value(BinaryKind kind) { + return static_cast(kind); +} + +static bool unary_kind_from_ggml_unary_op(ggml_unary_op op, UnaryKind & kind) { + switch (op) { + case GGML_UNARY_OP_ABS: + kind = UnaryKind::Abs; + return true; + case GGML_UNARY_OP_SGN: + kind = UnaryKind::Sgn; + return true; + case GGML_UNARY_OP_NEG: + kind = UnaryKind::Neg; + return true; + case GGML_UNARY_OP_STEP: + kind = UnaryKind::Step; + return true; + case GGML_UNARY_OP_TANH: + kind = UnaryKind::Tanh; + return true; + case GGML_UNARY_OP_ELU: + kind = UnaryKind::Elu; + return true; + case GGML_UNARY_OP_RELU: + kind = UnaryKind::Relu; + return true; + case GGML_UNARY_OP_SIGMOID: + kind = UnaryKind::Sigmoid; + return true; + case GGML_UNARY_OP_GELU: + kind = UnaryKind::Gelu; + return true; + case GGML_UNARY_OP_GELU_QUICK: + kind = UnaryKind::GeluQuick; + return true; + case GGML_UNARY_OP_SILU: + kind = UnaryKind::Silu; + return true; + case GGML_UNARY_OP_HARDSWISH: + kind = UnaryKind::HardSwish; + return true; + case GGML_UNARY_OP_HARDSIGMOID: + kind = UnaryKind::HardSigmoid; + return true; + case GGML_UNARY_OP_EXP: + kind = UnaryKind::Exp; + return true; + case GGML_UNARY_OP_EXPM1: + kind = UnaryKind::Expm1; + return true; + case GGML_UNARY_OP_SOFTPLUS: + kind = UnaryKind::SoftPlus; + return true; + case GGML_UNARY_OP_GELU_ERF: + kind = UnaryKind::GeluErf; + return true; + case GGML_UNARY_OP_XIELU: + kind = UnaryKind::Xielu; + return true; + case GGML_UNARY_OP_FLOOR: + kind = UnaryKind::Floor; + return true; + case GGML_UNARY_OP_CEIL: + kind = UnaryKind::Ceil; + return true; + case GGML_UNARY_OP_ROUND: + kind = UnaryKind::Round; + return true; + case GGML_UNARY_OP_TRUNC: + kind = UnaryKind::Trunc; + return true; + default: + return false; + } +} + +bool import_unary_kind(const ggml_tensor & tensor, UnaryKind & kind) { + switch (tensor.op) { + case GGML_OP_UNARY: + return unary_kind_from_ggml_unary_op(ggml_get_unary_op(&tensor), kind); + case GGML_OP_SQR: + kind = UnaryKind::Sqr; + return true; + case GGML_OP_SQRT: + kind = UnaryKind::Sqrt; + return true; + case GGML_OP_LOG: + kind = UnaryKind::Log; + return true; + case GGML_OP_SIN: + kind = UnaryKind::Sin; + return true; + case GGML_OP_COS: + kind = UnaryKind::Cos; + return true; + default: + return false; + } +} + +bool unary_kind_supported(UnaryKind kind) { + return static_cast(kind) <= static_cast(UnaryKind::Identity); +} + +uint32_t unary_kind_config_value(UnaryKind kind) { + return static_cast(kind); +} + OpParams import_op_params(const ggml_tensor & tensor) { + BinaryKind binary_kind; + if (import_binary_kind(tensor, binary_kind)) { + return BinaryParams{ binary_kind }; + } + + UnaryKind unary_kind; + if (import_unary_kind(tensor, unary_kind)) { + return UnaryParams{ unary_kind }; + } + switch (tensor.op) { case GGML_OP_RMS_NORM: + case GGML_OP_L2_NORM: return RmsNormParams{ ggml_get_op_params_f32(&tensor, 0) }; case GGML_OP_SOFT_MAX: return SoftMaxParams{ @@ -93,14 +279,29 @@ OpParams import_op_params(const ggml_tensor & tensor) { ggml_get_op_params_f32(&tensor, 1), }; case GGML_OP_GLU: - return GluParams{ ggml_get_glu_op(&tensor) }; + return GluParams{ ggml_get_glu_op(&tensor), ggml_get_op_params_i32(&tensor, 1) != 0 }; + case GGML_OP_SCALE: + return ScaleParams{ + ggml_get_op_params_f32(&tensor, 0), + ggml_get_op_params_f32(&tensor, 1), + }; case GGML_OP_ROPE: return RopeParams{ - ggml_get_op_params_i32(&tensor, 1), ggml_get_op_params_i32(&tensor, 2), - ggml_get_op_params_i32(&tensor, 4), ggml_get_op_params_f32(&tensor, 5), - ggml_get_op_params_f32(&tensor, 6), ggml_get_op_params_f32(&tensor, 7), - ggml_get_op_params_f32(&tensor, 8), ggml_get_op_params_f32(&tensor, 9), + ggml_get_op_params_i32(&tensor, 1), + ggml_get_op_params_i32(&tensor, 2), + ggml_get_op_params_i32(&tensor, 4), + ggml_get_op_params_f32(&tensor, 5), + ggml_get_op_params_f32(&tensor, 6), + ggml_get_op_params_f32(&tensor, 7), + ggml_get_op_params_f32(&tensor, 8), + ggml_get_op_params_f32(&tensor, 9), ggml_get_op_params_f32(&tensor, 10), + { + ggml_get_op_params_i32(&tensor, 11), + ggml_get_op_params_i32(&tensor, 12), + ggml_get_op_params_i32(&tensor, 13), + ggml_get_op_params_i32(&tensor, 14), + }, }; default: return std::monostate{}; @@ -110,6 +311,7 @@ OpParams import_op_params(const ggml_tensor & tensor) { bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs) { switch (op) { case GGML_OP_RMS_NORM: + case GGML_OP_L2_NORM: return rms_norm_params_equivalent(lhs, rhs); case GGML_OP_SOFT_MAX: return soft_max_params_equivalent(lhs, rhs); @@ -120,7 +322,21 @@ bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs case GGML_OP_CLAMP: return clamp_params_equivalent(lhs, rhs); case GGML_OP_GLU: - return glu_params_equivalent(lhs, rhs); + return binary_params_equivalent(lhs, rhs) || glu_params_equivalent(lhs, rhs); + case GGML_OP_SCALE: + return scale_params_equivalent(lhs, rhs); + case GGML_OP_ADD: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + return binary_params_equivalent(lhs, rhs); + case GGML_OP_UNARY: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_SIN: + case GGML_OP_COS: + return unary_params_equivalent(lhs, rhs); case GGML_OP_ROPE: return rope_params_equivalent(lhs, rhs); default: diff --git a/ggml/src/ggml-hrx/graph/op-params.h b/ggml/src/ggml-hrx/graph/op-params.h index dd1a97ebbdf8..fc6a7354f60f 100644 --- a/ggml/src/ggml-hrx/graph/op-params.h +++ b/ggml/src/ggml-hrx/graph/op-params.h @@ -2,6 +2,8 @@ #include "ggml.h" +#include +#include #include struct ggml_tensor; @@ -35,6 +37,65 @@ struct ClampParams { struct GluParams { ggml_glu_op op = GGML_GLU_OP_REGLU; + bool swapped = false; +}; + +struct ScaleParams { + float scale = 0.0f; + float bias = 0.0f; +}; + +enum class BinaryKind : uint32_t { + Add = 0, + Sub = 1, + Mul = 2, + Div = 3, + SwiGLU = 4, + GeGLU = 5, + RegLU = 6, + GeGLUErf = 7, + GeGLUQuick = 8, +}; + +struct BinaryParams { + BinaryKind op = BinaryKind::Add; +}; + +enum class UnaryKind : uint32_t { + Neg = 0, + Abs = 1, + Relu = 2, + Step = 3, + Sqr = 4, + Sgn = 5, + Floor = 6, + Ceil = 7, + Round = 8, + Trunc = 9, + Tanh = 10, + Elu = 11, + Sigmoid = 12, + Gelu = 13, + GeluQuick = 14, + Silu = 15, + HardSwish = 16, + HardSigmoid = 17, + Exp = 18, + Expm1 = 19, + GeluErf = 20, + Sqrt = 21, + Log = 22, + Identity = 23, + + // Parameterized or remaining exact-math unary routes need separate lowering work. + SoftPlus = 100, + Xielu, + Sin, + Cos, +}; + +struct UnaryParams { + UnaryKind op = UnaryKind::Abs; }; struct RopeParams { @@ -47,6 +108,7 @@ struct RopeParams { float attn_factor = 0.0f; float beta_fast = 0.0f; float beta_slow = 0.0f; + std::array sections = {}; }; // clang-format off @@ -58,6 +120,9 @@ using OpParams = std::variant< ArgsortParams, ClampParams, GluParams, + ScaleParams, + BinaryParams, + UnaryParams, RopeParams>; // clang-format on @@ -68,5 +133,11 @@ template const T * op_params_as(const OpParams & params) { OpParams import_op_params(const ggml_tensor & tensor); bool op_params_equivalent(ggml_op op, const OpParams & lhs, const OpParams & rhs); bool op_params_equivalent(ggml_op op, const OpParams & lhs, const ggml_tensor & rhs); +bool import_binary_kind(const ggml_tensor & tensor, BinaryKind & kind); +bool binary_kind_supported(BinaryKind kind); +uint32_t binary_kind_config_value(BinaryKind kind); +bool import_unary_kind(const ggml_tensor & tensor, UnaryKind & kind); +bool unary_kind_supported(UnaryKind kind); +uint32_t unary_kind_config_value(UnaryKind kind); } // namespace ggml::hrx diff --git a/ggml/src/ggml-hrx/graph/value-map.cpp b/ggml/src/ggml-hrx/graph/value-map.cpp index 8bc09d186877..dfeea7c027ab 100644 --- a/ggml/src/ggml-hrx/graph/value-map.cpp +++ b/ggml/src/ggml-hrx/graph/value-map.cpp @@ -5,12 +5,66 @@ #include namespace ggml::hrx { +namespace { + +static const ggml_tensor * tensor_storage_root(const ggml_tensor * tensor) { + while (tensor != nullptr && tensor->view_src != nullptr) { + tensor = tensor->view_src; + } + return tensor; +} + +static size_t tensor_storage_offset(const ggml_tensor * tensor) { + return tensor != nullptr && tensor->view_src != nullptr ? tensor->view_offs : 0; +} + +static bool tensor_storage_relative_offset(const ggml_tensor * source, const ggml_tensor * tensor, size_t & offset) { + if (source == nullptr || tensor == nullptr || tensor_storage_root(source) != tensor_storage_root(tensor)) { + return false; + } + const size_t source_offset = tensor_storage_offset(source); + const size_t tensor_offset = tensor_storage_offset(tensor); + if (tensor_offset < source_offset) { + return false; + } + const size_t relative_offset = tensor_offset - source_offset; + if (relative_offset > ggml_nbytes(source) || ggml_nbytes(tensor) > ggml_nbytes(source) - relative_offset) { + return false; + } + offset = relative_offset; + return true; +} + +} // namespace const Value * ValueMap::find_alias_source(const ggml_tensor * tensor) const { if (tensor == nullptr || tensor->view_src == nullptr) { return nullptr; } - return find_tensor(tensor->view_src); + const Value * exact_source = find_tensor(tensor->view_src); + if (exact_source != nullptr) { + return exact_source; + } + const Value * best_source = nullptr; + size_t best_source_size = 0; + for (const Value & value : values_) { + size_t relative_offset = 0; + if (value.tensor == nullptr || !tensor_storage_relative_offset(value.tensor, tensor, relative_offset)) { + continue; + } + if (best_source == nullptr || ggml_nbytes(value.tensor) < best_source_size) { + best_source = &value; + best_source_size = ggml_nbytes(value.tensor); + } + } + return best_source; +} + +void ValueMap::promote_storage_root_external(ValueId storage_root) { + if (storage_root.value < 0 || static_cast(storage_root.value) >= values_.size()) { + return; + } + values_[static_cast(storage_root.value)].kind = ValueKind::External; } ValueId ValueMap::get_or_add_tensor_value(const ggml_tensor * tensor, ValueKind kind) { @@ -19,6 +73,7 @@ ValueId ValueMap::get_or_add_tensor_value(const ggml_tensor * tensor, ValueKind Value & value = values_[found->second]; if (kind == ValueKind::External) { value.kind = ValueKind::External; + promote_storage_root_external(value.storage_root); } return value.id; } @@ -31,15 +86,22 @@ ValueId ValueMap::get_or_add_tensor_value(const ggml_tensor * tensor, ValueKind size_t storage_offset = 0; size_t storage_byte_count = ggml_nbytes(tensor); if (alias_source != nullptr) { + size_t relative_offset = 0; + if (!tensor_storage_relative_offset(alias_source->tensor, tensor, relative_offset)) { + relative_offset = tensor->view_offs; + } storage = alias_source->storage; storage_root = alias_source->storage_root; alias_source_id = alias_source->id; - storage_offset = tensor->view_offs; + storage_offset = alias_source->storage_offset + relative_offset; storage_byte_count = alias_source->storage_byte_count; const Value * root = find(storage_root); if (root != nullptr && root->kind == ValueKind::External) { kind = ValueKind::External; } + if (kind == ValueKind::External) { + promote_storage_root_external(storage_root); + } } else { storage = ValueStorageId(static_cast(storages_.size())); storage_root = id; diff --git a/ggml/src/ggml-hrx/graph/value-map.h b/ggml/src/ggml-hrx/graph/value-map.h index 783ffd0519ea..b026de65fde6 100644 --- a/ggml/src/ggml-hrx/graph/value-map.h +++ b/ggml/src/ggml-hrx/graph/value-map.h @@ -118,6 +118,7 @@ class ValueMap { private: const Value * find_alias_source(const ggml_tensor * tensor) const; + void promote_storage_root_external(ValueId storage_root); std::vector values_; std::vector storages_; diff --git a/ggml/src/ggml-hrx/hrx-minimum-commit.txt b/ggml/src/ggml-hrx/hrx-minimum-commit.txt new file mode 100644 index 000000000000..59dce28e2605 --- /dev/null +++ b/ggml/src/ggml-hrx/hrx-minimum-commit.txt @@ -0,0 +1 @@ +32d0c76a600f9c8009bb61e25f594570f2b0658a diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/gather_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom similarity index 69% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/gather_add_f32.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom index 6bef54908bc7..c8e46116e5d6 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/gather_add_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/gather_add_f32.loom @@ -9,19 +9,20 @@ // GGML GET_ROWS operations followed by ADD. amdgpu.target @ggml_gather_add_gfx11_wave64 {subgroup_size = 64} -kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ggml_gather_add_f32( - %source_token_count: index, %output_token_count: index, %hidden_size: index) { +config.decl @ggml.gather_add_f32.binary_op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ggml_gather_add_f32(%source_token_count: index, %output_token_count: index, %hidden_size: index) { %one = index.constant 1 : index %twofiftysix = index.constant 256 : index %rounding = index.constant 255 : index %rounded_width = index.add %hidden_size, %rounding : index %column_workgroup_count = index.div %rounded_width, %twofiftysix : index kernel.launch.config workgroups(%column_workgroup_count, %output_token_count, %one) workgroup_size(%twofiftysix, %one, %one) : index -} launch(%source_token_count: index, %output_token_count: index, %hidden_size: index, - %attention: buffer, %residual: buffer, %output_ids: buffer, %output: buffer) { +} launch(%source_token_count: index, %output_token_count: index, %hidden_size: index, %attention: buffer, %residual: buffer, %output_ids: buffer, %output: buffer) { %source_count = index.assume %source_token_count [range(%source_token_count, 1, 2048)] : index %output_count = index.assume %output_token_count [range(%output_token_count, 1, 2048)] : index %width = index.assume %hidden_size [range(%hidden_size, 1, 65536)] : index + %binary_op = config.get @ggml.gather_add_f32.binary_op : index %column_workgroup = kernel.workgroup.id : index %output_row0 = kernel.workgroup.id : index %workitem = kernel.workitem.id : index @@ -45,7 +46,19 @@ kernel.def target(@ggml_gather_add_gfx11_wave64) export("ggml_gather_add_f32") @ %attention_value = view.load %attention_view[%source_row, %safe_column] : view<[%source_count]x[%width]xf32> -> f32 %residual_value = view.load %residual_view[%source_row, %safe_column] : view<[%source_count]x[%width]xf32> -> f32 %sum = scalar.addf %attention_value, %residual_value : f32 - view.store %sum, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%width]xf32> + %difference = scalar.subf %attention_value, %residual_value : f32 + %product = scalar.mulf %attention_value, %residual_value : f32 + %quotient = scalar.divf %attention_value, %residual_value : f32 + %op0 = index.constant 0 : index + %op1 = index.constant 1 : index + %op2 = index.constant 2 : index + %is_add = index.cmp eq, %binary_op, %op0 : index + %is_sub = index.cmp eq, %binary_op, %op1 : index + %is_mul = index.cmp eq, %binary_op, %op2 : index + %add_or_div = scf.select %is_add, %sum, %quotient : f32 + %sub_or_other = scf.select %is_sub, %difference, %add_or_div : f32 + %result = scf.select %is_mul, %product, %sub_or_other : f32 + view.store %result, %output_view[%output_row, %safe_column] : f32, view<[%launch_output_count]x[%width]xf32> } kernel.return } @@ -55,12 +68,12 @@ check.case public @ggml_gather_add_noncontiguous_case { %one = check.literal value(1) : index %three = check.literal value(3) : index %four = check.literal value(4) : index - %attention = check.generate.iota offset(0) step(1) : tensor<3x4xf32> - %residual = check.generate.fill value(100) : tensor<3x4xf32> + %attention = check.generate.iota offset(0.0) step(1.0) : tensor<3x4xf32> + %residual = check.generate.fill value(100.0) : tensor<3x4xf32> %output_ids = check.generate.fill value(2) : tensor<1xi32> - %output = check.generate.fill value(0) : tensor<1x4xf32> - %expected = check.generate.iota offset(108) step(1) : tensor<1x4xf32> - func.call @ggml_gather_add_f32(%three, %one, %four, %attention, %residual, %output_ids, %output) : (index, index, index, tensor<3x4xf32>, tensor<3x4xf32>, tensor<1xi32>, tensor<1x4xf32>) + %output = check.generate.fill value(0.0) : tensor<1x4xf32> + %expected = check.generate.iota offset(108.0) step(1.0) : tensor<1x4xf32> + kernel.launch @ggml_gather_add_f32[%three, %one, %four](%three, %one, %four, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<3x4xf32>, tensor<3x4xf32>, tensor<1xi32>, tensor<1x4xf32>) check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x4xf32> check.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/manifest.json new file mode 100644 index 000000000000..572b6756e1d9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx/manifest.json @@ -0,0 +1,68 @@ +{ + "schema": "ggml-hrx-kernel-corpus-v1", + "upstream_revision": "local", + "files": [ + { + "path": "gather_add_f32.loom" + } + ], + "exports": [ + { + "name": "ggml_gather_add_f32", + "symbol": "ggml_gather_add_f32", + "source": "gather_add_f32.loom", + "workload_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "attention", + "residual", + "output_ids", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "gather_add_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [], + "family": "hrx" + } + ], + "link_modules": [], + "plan_cases": [] +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/add_f32.loom deleted file mode 100644 index 1797b01935c7..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/hrx_owned/add_f32.loom +++ /dev/null @@ -1,47 +0,0 @@ -// Copyright 2026 The HRX Authors -// SPDX-License-Identifier: Apache-2.0 - -amdgpu.target @ggml_add_gfx11_wave64 {subgroup_size = 64} - -kernel.def target(@ggml_add_gfx11_wave64) export("ggml_add_f32") @ggml_add_f32(%element_count: index) { - %one = index.constant 1 : index - %twofiftysix = index.constant 256 : index - %rounding = index.constant 255 : index - %rounded = index.add %element_count, %rounding : index - %workgroup_count = index.div %rounded, %twofiftysix : index - kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index -} launch(%element_count: index, %a: buffer, %b: buffer, %output: buffer) { - %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index - %workgroup = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %twofiftysix = index.constant 256 : index - %base0 = index.mul %workgroup, %twofiftysix : index - %linear0 = index.add %base0, %workitem : index - %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index - %in_bounds = index.cmp ult, %linear, %count : index - %zero_offset = index.constant 0 : offset - %a_noalias, %b_noalias, %output_noalias = buffer.assume.noalias %a, %b, %output : buffer, buffer, buffer - %a_view = buffer.view %a_noalias[%zero_offset] : buffer -> view<[%count]xf32> - %b_view = buffer.view %b_noalias[%zero_offset] : buffer -> view<[%count]xf32> - %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%count]xf32> - scf.if %in_bounds { - %a_value = view.load %a_view[%linear] : view<[%count]xf32> -> f32 - %b_value = view.load %b_view[%linear] : view<[%count]xf32> -> f32 - %sum = scalar.addf %a_value, %b_value : f32 - view.store %sum, %output_view[%linear] : f32, view<[%count]xf32> - } - kernel.return -} - -check.case public @ggml_add_f32_small_case { - %four = check.literal value(4) : index - %a = check.generate.iota offset(1) step(1) : tensor<4xf32> - %b = check.generate.fill value(10) : tensor<4xf32> - %output = check.generate.fill value(0) : tensor<4xf32> - %expected = check.generate.iota offset(11) step(1) : tensor<4xf32> - func.call @ggml_add_f32(%four, %a, %b, %output) : (index, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> - check.return -} - -check.benchmark<@ggml_add_f32_small_case> @ggml_add_f32_small 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 new file mode 100644 index 000000000000..625075d1fbc9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/manifest.json @@ -0,0 +1,9645 @@ +{ + "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" + }, + { + "path": "ops/elementwise/binary_swiglu_symmetric_i4_k32.loom" + }, + { + "path": "ops/elementwise/binary_bc_f32.loom" + }, + { + "path": "ops/elementwise/binary_bc_publish_q8_1_x4.loom" + }, + { + "path": "ops/elementwise/scale_f32.loom" + }, + { + "path": "ops/elementwise/scale_add_f32.loom" + }, + { + "path": "ops/elementwise/scale_bias_f32.loom" + }, + { + "path": "ops/copy/copy_f32.loom" + }, + { + "path": "ops/copy/copy_f32_f16.loom" + }, + { + "path": "ops/copy/copy_f16_k16_major.loom" + }, + { + "path": "ops/copy/transpose_f16.loom" + }, + { + "path": "ops/copy/strided_source_f32.loom" + }, + { + "path": "ops/copy/concat_dim0_f32.loom" + }, + { + "path": "ops/copy/concat_dim0_strided_source_f32.loom" + }, + { + "path": "ops/elementwise/binary_f32.loom" + }, + { + "path": "ops/elementwise/binary_publish_q8_1_x4.loom" + }, + { + "path": "ops/normalization/rmsnorm_f32.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_f32.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_strided_f32.loom" + }, + { + "path": "ops/normalization/rmsnorm_mul_add_f32.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_f32_f16.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_f32_k16.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom" + }, + { + "path": "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + }, + { + "path": "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + }, + { + "path": "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + }, + { + "path": "ops/normalization/rmsnorm_gate_publish.loom" + }, + { + "path": "ops/normalization/rmsnorm_mul_rope_f32.loom" + }, + { + "path": "ops/mul_mat_f32_f32_wmma.loom" + }, + { + "path": "ops/matmul/vector_f32_f32.loom" + }, + { + "path": "ops/matmul/vector_postops_f32_f32.loom" + }, + { + "path": "motifs/matmul/vector_publish.loom" + }, + { + "path": "motifs/matmul/vector_decode_pair.loom" + }, + { + "path": "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom" + }, + { + "path": "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom" + }, + { + "path": "ops/matmul_id/mul_mat_id_f16_f16_wmma.loom" + }, + { + "path": "ops/matmul_id/mul_mat_id_swiglu_f16_f16_wmma.loom" + }, + { + "path": "motifs/matmul_id/f16_pair_schedule.loom" + }, + { + "path": "motifs/matmul_id/f16_pair_accumulate.loom" + }, + { + "path": "motifs/matmul_id/f16_pair_binary_publish.loom" + }, + { + "path": "ops/matmul_id/tiled_input_f32_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_input_f32_publish_f32.loom" + }, + { + "path": "ops/matmul_id/tiled_input_f32_postops_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_input_f32_postops_publish_f32.loom" + }, + { + "path": "ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom" + }, + { + "path": "ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom" + }, + { + "path": "ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom" + }, + { + "path": "ops/matmul_id/vector_pair_binary.loom" + }, + { + "path": "ops/matmul_id/vector_weighted.loom" + }, + { + "path": "motifs/normalization/rmsnorm_q8_1_x4.loom" + }, + { + "path": "motifs/matmul/q4_q8_1_x4_decode.loom" + }, + { + "path": "motifs/matmul/q4_q8_1_x4_dense.loom" + }, + { + "path": "motifs/matmul/q4_q6_q8_1_x4_qkv.loom" + }, + { + "path": "motifs/matmul/q6_f32_decode.loom" + }, + { + "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" + }, + { + "path": "motifs/routing/router_topk_row.loom" + }, + { + "path": "motifs/routing/router_projection_fused.loom" + }, + { + "path": "motifs/routing/weighted_reduce_f16_f32.loom" + }, + { + "path": "motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" + }, + { + "path": "motifs/attention/rmsnorm_neox_cache_f16.loom" + }, + { + "path": "motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul_id/pair_swiglu_publish_f32.loom" + }, + { + "path": "motifs/matmul_id/vector_q4k_q8.loom" + }, + { + "path": "motifs/matmul_id/vector_pair.loom" + }, + { + "path": "motifs/matmul_id/vector_weighted.loom" + }, + { + "path": "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul_id/tiled_postprocess_f32.loom" + }, + { + "path": "motifs/matmul_id/skinny_postprocess_f32.loom" + }, + { + "path": "ops/moe_routing_tables.loom" + }, + { + "path": "ops/mul_mat_swiglu_f32_f32_wmma.loom" + }, + { + "path": "ops/mul_mat_swiglu_f32_f32_decode.loom" + }, + { + "path": "ops/mul_mat_add_f32_f32_decode.loom" + }, + { + "path": "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + }, + { + "path": "ops/mul_mat_symmetric_i4_wmma.loom" + }, + { + "path": "ops/mul_mat_symmetric_i8_wmma.loom" + }, + { + "path": "ops/mul_mat_q5_k_q8_plane_wmma.loom" + }, + { + "path": "ops/mul_mat_q6_k_packed_f16_wmma.loom" + }, + { + "path": "ops/matmul/vector_q6_f32_f32_kernel.loom" + }, + { + "path": "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom" + }, + { + "path": "ops/matmul/skinny_input_f32_publish_f32.loom" + }, + { + "path": "ops/mul_mat_dual_q4_f32_decode.loom" + }, + { + "path": "ops/mul_mat_f32_f32_decode.loom" + }, + { + "path": "ops/rows/get_rows_f32.loom" + }, + { + "path": "ops/rows/get_rows_scale_f32.loom" + }, + { + "path": "ops/rows/gather_add_rmsnorm_binary_f32.loom" + }, + { + "path": "ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom" + }, + { + "path": "ops/rows/get_rows_next_f32.loom" + }, + { + "path": "ops/gated_delta_net/head128_scan_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_scan_projection_epilogue_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_scan_rmsnorm_gate_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_inplace_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_inplace_projection_epilogue_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_snapshot_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_snapshot_projection_epilogue_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_selected_snapshot_projection_epilogue_f32_wmma.loom" + }, + { + "path": "ops/gated_delta_net/head128_selected_snapshot_projection_rms_gate_publish_q8.loom" + }, + { + "path": "motifs/gated_delta_net/scan_state_head128.loom" + }, + { + "path": "motifs/gated_delta_net/rmsnorm_gate.loom" + }, + { + "path": "motifs/gated_delta_net/inplace_state_head128.loom" + }, + { + "path": "motifs/gated_delta_net/snapshot_load.loom" + }, + { + "path": "motifs/gated_delta_net/snapshot_publish.loom" + }, + { + "path": "motifs/gated_delta_net/selected_snapshot_load.loom" + }, + { + "path": "motifs/gated_delta_net/publish_f32.loom" + }, + { + "path": "motifs/gated_delta_net/publish_q8.loom" + }, + { + "path": "ops/gated_delta_net/projection_epilogue_f32.loom" + }, + { + "path": "motifs/gated_delta_net/projection_epilogue.loom" + }, + { + "path": "motifs/ssm_conv/generic_apply.loom" + }, + { + "path": "motifs/ssm_conv/dconv4_silu.loom" + }, + { + "path": "motifs/ssm_conv/prefill_finish.loom" + }, + { + "path": "motifs/ssm_conv/snapshot_window_tail.loom" + }, + { + "path": "motifs/ssm_conv/prefill_window.loom" + }, + { + "path": "motifs/ssm_conv/rollback_store.loom" + }, + { + "path": "motifs/ssm_conv/decode_update.loom" + }, + { + "path": "motifs/ssm_conv/state_materialize.loom" + }, + { + "path": "ops/ssm_conv/generic_f32.loom" + }, + { + "path": "ops/ssm_conv/lfm_dconv3_state_f32.loom" + }, + { + "path": "ops/ssm_conv/generic_binary_f32.loom" + }, + { + "path": "ops/ssm_conv/prefill_finish_f32.loom" + }, + { + "path": "ops/ssm_conv/snapshot_window_tail_f32.loom" + }, + { + "path": "ops/ssm_conv/prefill_dconv4_silu_512_wg1024.loom" + }, + { + "path": "ops/ssm_conv/decode_dconv4_silu_f32.loom" + }, + { + "path": "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + }, + { + "path": "ops/rope/rope_f32.loom" + }, + { + "path": "ops/rope/rope_concat_f32.loom" + }, + { + "path": "ops/rope/rope_scale_f32.loom" + }, + { + "path": "ops/rope/rope_set_rows_f32.loom" + }, + { + "path": "ops/rows/set_rows.loom" + }, + { + "path": "ops/flash_attention/prefill_f32_f16_wmma.loom" + }, + { + "path": "ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom" + }, + { + "path": "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom" + }, + { + "path": "ops/flash_attention/decode_split_next_q8_test.loom" + }, + { + "path": "motifs/flash_attention/qk_block_wmma.loom" + }, + { + "path": "motifs/flash_attention/online_softmax.loom" + }, + { + "path": "motifs/flash_attention/pv_block_wmma.loom" + }, + { + "path": "motifs/flash_attention/publish_f32.loom" + }, + { + "path": "motifs/flash_attention/decode_partials.loom" + }, + { + "path": "motifs/flash_attention/decode_reduce.loom" + }, + { + "path": "motifs/flash_attention/publish_q8.loom" + }, + { + "path": "motifs/flash_attention/completion_counter_reduce.loom" + }, + { + "path": "motifs/elementwise/unary_apply.loom" + }, + { + "path": "motifs/elementwise/symmetric_i4_binary.loom" + }, + { + "path": "motifs/elementwise/broadcast_index.loom" + }, + { + "path": "motifs/elementwise/scale_apply.loom" + }, + { + "path": "motifs/elementwise/binary_apply.loom" + }, + { + "path": "motifs/copy/vector_load_store.loom" + }, + { + "path": "motifs/copy/convert_f32_f16.loom" + }, + { + "path": "motifs/copy/layout_k16_major.loom" + }, + { + "path": "motifs/copy/transpose_tile.loom" + }, + { + "path": "motifs/copy/concat_dim0.loom" + }, + { + "path": "motifs/copy/strided_source.loom" + }, + { + "path": "motifs/normalization/rmsnorm.loom" + }, + { + "path": "motifs/normalization/rmsnorm_binary.loom" + }, + { + "path": "motifs/normalization/rmsnorm_symmetric_i4.loom" + }, + { + "path": "motifs/normalization/rmsnorm_gate.loom" + }, + { + "path": "motifs/normalization/publish_f16.loom" + }, + { + "path": "motifs/normalization/publish_q8.loom" + }, + { + "path": "motifs/quantize_q8_1_x4.loom" + }, + { + "path": "motifs/normalization/rmsnorm_q8_1_x4.loom" + }, + { + "path": "motifs/quantize_q8_1_x4_group.loom" + }, + { + "path": "motifs/publish_f32.loom" + }, + { + "path": "motifs/rope/rope_f32.loom" + }, + { + "path": "motifs/rows/cache_publish.loom" + }, + { + "path": "motifs/rows/get_rows_decode.loom" + }, + { + "path": "motifs/dequant.loom" + }, + { + "path": "motifs/mul_mat_f32_f32_wmma_core.loom" + }, + { + "path": "motifs/matmul/vector_activation.loom" + }, + { + "path": "motifs/matmul/vector_postops.loom" + }, + { + "path": "motifs/matmul/vector_q6_f32_body.loom" + }, + { + "path": "motifs/matmul/vector_publish_f32.loom" + }, + { + "path": "motifs/mul_mat_quantized_f16_prefill.loom" + }, + { + "path": "motifs/matmul_id/f16_down_schedule.loom" + }, + { + "path": "motifs/matmul_id/f16_publish.loom" + }, + { + "path": "motifs/matmul_id/f16_projection_q4k.loom" + }, + { + "path": "motifs/matmul_id/f16_projection_q5k.loom" + }, + { + "path": "motifs/matmul_id/f16_projection_q6k.loom" + }, + { + "path": "motifs/mul_mat_f32_f32_postops.loom" + }, + { + "path": "motifs/llm_attention_qkv_matmul_postops.loom" + }, + { + "path": "motifs/q4_k_f16.loom" + }, + { + "path": "motifs/q6_k_f16.loom" + }, + { + "path": "motifs/q8_0_f16.loom" + }, + { + "path": "motifs/q8_1_f16.loom" + }, + { + "path": "motifs/f16_f16.loom" + }, + { + "path": "motifs/bf16_f16.loom" + }, + { + "path": "motifs/f32_f16.loom" + }, + { + "path": "motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul/skinny_postprocess_f32.loom" + }, + { + "path": "motifs/matmul/skinny_publish_f32.loom" + }, + { + "path": "ops/matmul/tiled_input_f32_publish_f32.loom" + }, + { + "path": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + }, + { + "path": "ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom" + }, + { + "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" + }, + { + "path": "motifs/matmul/tiled_publish.loom" + }, + { + "path": "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom" + }, + { + "path": "motifs/matmul/tiled_pair_postprocess_f32.loom" + } + ], + "exports": [ + { + "name": "ggml_mul_mat_vector_iq2_s_f32_f32", + "family": "loom_libs", + "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", + "weight", + "iq2_s_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_vector_iq3_s_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_iq3_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", + "weight", + "iq3_s_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_iq2_s_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_iq2_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", + "iq2_s_grid", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" + ], + "library_sources": [ + "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/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_mul_mat_tiled_input_f32_iq3_s_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_input_f32_iq3_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", + "iq3_s_grid", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" + ], + "library_sources": [ + "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/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_transpose_f16", + "family": "loom_libs", + "symbol": "ggml_copy_transpose_f16", + "source": "ops/copy/transpose_f16.loom", + "workload_parameters": [ + { + "name": "row_count", + "type": "index" + }, + { + "name": "column_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "row_count", + "type": "index" + }, + { + "name": "column_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/transpose_f16.loom" + ], + "library_sources": [ + "motifs/copy/transpose_tile.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/transpose_tile.loom" + ] + }, + { + "name": "ggml_copy_f16_k16_major", + "family": "loom_libs", + "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": [ + "input", + "output" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/copy/copy_f16_k16_major.loom" + ], + "library_sources": [ + "motifs/copy/layout_k16_major.loom", + "motifs/copy/convert_f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/layout_k16_major.loom", + "motifs/copy/convert_f32_f16.loom" + ] + }, + { + "name": "ggml_copy_f32_f16", + "family": "loom_libs", + "symbol": "ggml_copy_f32_f16", + "source": "ops/copy/copy_f32_f16.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "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" + }, + { + "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_strided_source_f32.loom" + ], + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ] + }, + "compile_dependencies": [ + "motifs/copy/concat_dim0.loom" + ] + }, + { + "name": "ggml_scale_bias_f32", + "family": "loom_libs", + "symbol": "ggml_scale_bias_f32", + "source": "ops/elementwise/scale_bias_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/scale_bias_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/scale_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/scale_apply.loom" + ] + }, + { + "name": "ggml_unary_f32", + "family": "loom_libs", + "symbol": "ggml_unary_f32", + "source": "ops/elementwise/unary_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/unary_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/unary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/unary_apply.loom" + ] + }, + { + "name": "ggml_scale_f32", + "family": "loom_libs", + "symbol": "ggml_scale_f32", + "source": "ops/elementwise/scale_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/scale_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/scale_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/scale_apply.loom" + ] + }, + { + "name": "ggml_scale_add_f32", + "family": "loom_libs", + "symbol": "ggml_scale_add_f32", + "source": "ops/elementwise/scale_add_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "input", + "residual", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/scale_add_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/scale_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/scale_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_binary_f32", + "family": "loom_libs", + "symbol": "ggml_binary_f32", + "source": "ops/elementwise/binary_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/binary_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "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", + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "source": "ops/elementwise/binary_swiglu_symmetric_i4_k32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "lhs", + "rhs", + "output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/binary_swiglu_symmetric_i4_k32.loom" + ], + "library_sources": [ + "motifs/elementwise/symmetric_i4_binary.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/symmetric_i4_binary.loom" + ] + }, + { + "name": "ggml_binary_bc_f32", + "family": "loom_libs", + "symbol": "ggml_binary_bc_f32", + "source": "ops/elementwise/binary_bc_f32.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + }, + { + "name": "ne0", + "type": "index" + }, + { + "name": "ne1", + "type": "index" + }, + { + "name": "ne2", + "type": "index" + }, + { + "name": "ne3", + "type": "index" + }, + { + "name": "src0_element_count", + "type": "index" + }, + { + "name": "src1_element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + }, + { + "name": "ne0", + "type": "index" + }, + { + "name": "ne1", + "type": "index" + }, + { + "name": "ne2", + "type": "index" + }, + { + "name": "ne3", + "type": "index" + }, + { + "name": "src0_element_count", + "type": "index" + }, + { + "name": "src1_element_count", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/elementwise/binary_bc_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/binary_apply.loom", + "motifs/elementwise/broadcast_index.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/binary_apply.loom", + "motifs/elementwise/broadcast_index.loom" + ] + }, + { + "name": "ggml_binary_bc_f32_publish_q8_1_x4", + "family": "loom_libs", + "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": [ + "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", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom" + ] + }, + { + "name": "ggml_rmsnorm_mul_rope_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_mul_rope_f32", + "source": "ops/normalization/rmsnorm_mul_rope_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "weight", + "positions", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_rmsnorm_binary_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_f32", + "source": "ops/normalization/rmsnorm_binary_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_binary_f32_lowtoken_striped", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_f32_lowtoken_striped", + "source": "ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_binary_strided_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_strided_f32", + "source": "ops/normalization/rmsnorm_binary_strided_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_strided_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_mul_add_f32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_mul_add_f32", + "source": "ops/normalization/rmsnorm_mul_add_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "residual", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_mul_add_f32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_binary_f32_k16", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_f32_k16", + "source": "ops/normalization/rmsnorm_binary_f32_k16.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32_k16.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/publish_f16.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/normalization/publish_f16.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_binary_f32_f16", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_f32_f16", + "source": "ops/normalization/rmsnorm_binary_f32_f16.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32_f16.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_rmsnorm_binary_q8_1_x4_publish", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_q8_1_x4_publish", + "source": "ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "rhs", + "output", + "q8_output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.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_symmetric_i4_k32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "source": "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ] + }, + { + "name": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "family": "loom_libs", + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "source": "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "lhs", + "rhs", + "residual_output", + "weight", + "output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "read", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ] + }, + { + "name": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "source": "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "raw_gate", + "output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ] + }, + "compile_dependencies": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ] + }, + { + "name": "ggml_rmsnorm_gate_f32_publish", + "family": "loom_libs", + "symbol": "ggml_rmsnorm_gate_f32_publish", + "source": "ops/normalization/rmsnorm_gate_publish.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "raw_gate", + "output", + "alternate_output" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/normalization/rmsnorm_gate_publish.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_mul_mat_quantized_f16_wmma_prefill_conv4_interior", + "family": "loom_libs", + "symbol": "ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior", + "source": "ops/mul_mat_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "filter", + "output", + "edges" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_quantized_f16_wmma_prefill_conv4", + "family": "loom_libs", + "symbol": "ggml_mul_mat_quantized_f16_wmma_prefill_conv4", + "source": "ops/mul_mat_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "state", + "filter", + "output", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_q4_k_f16_wmma_prefill_wave32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q4_k_f16_wmma_prefill_wave32", + "source": "ops/mul_mat_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_f32_f32_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_f32_f32_wmma", + "source": "ops/mul_mat_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "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", + "symbol": "ggml_mul_mat_vector_f32_f32", + "source": "ops/matmul/vector_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/dequant.loom", + "motifs/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_activation.loom", + "motifs/matmul/vector_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_vector_bias_residual_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_bias_residual_f32_f32", + "source": "ops/matmul/vector_postops_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "bias", + "residual_input", + "residual_output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_postops_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/matmul/vector_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_q_matmul_rope_tiled_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_q_matmul_rope_tiled_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "indices", + "theta", + "freq_factors", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_v_matmul_set_rows_tiled_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_v_matmul_set_rows_tiled_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "indices", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_q_matmul_rope_vector_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_q_matmul_rope_vector_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_k_matmul_rope_set_rows_vector_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_k_matmul_rope_set_rows_vector_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "indices", + "theta", + "freq_factors", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_v_matmul_set_rows_vector_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_v_matmul_set_rows_vector_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "indices", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom" + ], + "library_sources": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_activation.loom", + "motifs/matmul/vector_decode_pair.loom", + "motifs/matmul/vector_publish.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "llm_attention_q_matmul_rope_legacy_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_q_matmul_rope_legacy_f32_f32", + "source": "ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_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/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_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": "llm_attention_k_matmul_rope_set_rows_legacy_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_k_matmul_rope_set_rows_legacy_f32_f32", + "source": "ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "positions", + "indices", + "theta", + "freq_factors", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_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/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/rope/rope_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": "llm_attention_v_matmul_set_rows_legacy_f32_f32", + "family": "loom_libs", + "symbol": "llm_attention_v_matmul_set_rows_legacy_f32_f32", + "source": "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "indices", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.loom", + "motifs/dequant.loom", + "motifs/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/mul_mat_f32_f32_wmma_core.loom", + "motifs/llm_attention_qkv_matmul_postops.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_id_f16_f16_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_f16_f16_wmma", + "source": "ops/matmul_id/mul_mat_id_f16_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/mul_mat_id_f16_f16_wmma.loom" + ], + "library_sources": [ + "motifs/matmul_id/f16_down_schedule.loom", + "motifs/matmul_id/f16_publish.loom", + "motifs/matmul_id/f16_projection_q4k.loom", + "motifs/matmul_id/f16_projection_q6k.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/f16_down_schedule.loom", + "motifs/matmul_id/f16_publish.loom", + "motifs/matmul_id/f16_projection_q4k.loom", + "motifs/matmul_id/f16_projection_q6k.loom" + ] + }, + { + "name": "ggml_mul_mat_id_swiglu_f16_f16_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_swiglu_f16_f16_wmma", + "source": "ops/matmul_id/mul_mat_id_swiglu_f16_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/mul_mat_id_swiglu_f16_f16_wmma.loom" + ], + "library_sources": [ + "motifs/matmul_id/f16_pair_schedule.loom", + "motifs/matmul_id/f16_pair_accumulate.loom", + "motifs/matmul_id/f16_pair_binary_publish.loom", + "motifs/matmul_id/f16_projection_q4k.loom", + "motifs/matmul_id/f16_projection_q5k.loom", + "motifs/matmul_id/f16_projection_q6k.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/f16_pair_schedule.loom", + "motifs/matmul_id/f16_pair_accumulate.loom", + "motifs/matmul_id/f16_pair_binary_publish.loom", + "motifs/matmul_id/f16_projection_q4k.loom", + "motifs/matmul_id/f16_projection_q5k.loom", + "motifs/matmul_id/f16_projection_q6k.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_tiled_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_tiled_input_f32_publish_f32", + "source": "ops/matmul_id/tiled_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/tiled_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_skinny_input_f32_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_skinny_input_f32_publish_f32", + "source": "ops/matmul_id/skinny_input_f32_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/skinny_input_f32_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_tiled_input_f32_postops_publish_f32", + "source": "ops/matmul_id/tiled_input_f32_postops_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "bias", + "residual_input", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/tiled_input_f32_postops_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_skinny_input_f32_postops_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_skinny_input_f32_postops_publish_f32", + "source": "ops/matmul_id/skinny_input_f32_postops_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "bias", + "residual_input", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/skinny_input_f32_postops_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32", + "source": "ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "bias", + "residual_input", + "residual_output", + "norm_weight", + "normalized_output", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/tiled_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32", + "source": "ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "weight", + "bias", + "residual_input", + "residual_output", + "norm_weight", + "normalized_output", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/skinny_postprocess_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32", + "source": "ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/pair_swiglu_publish_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/pair_swiglu_publish_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32", + "source": "ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "expert_table", + "partition_table", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/pair_swiglu_publish_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul_id/pair_swiglu_publish_f32.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_id_vector_pair_binary_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_vector_pair_binary_publish_f32", + "source": "ops/matmul_id/vector_pair_binary.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_pair_binary.loom" + ], + "library_sources": [ + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_id_vector_pair_binary_publish_q8", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_vector_pair_binary_publish_q8", + "source": "ops/matmul_id/vector_pair_binary.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "gate_weight", + "up_weight", + "output", + "completion_counters", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_pair_binary.loom" + ], + "library_sources": [ + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/vector_pair.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_id_vector_weighted_wave32_publish_q8", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_vector_weighted_wave32_publish_q8", + "source": "ops/matmul_id/vector_weighted.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "route_ids", + "route_weights", + "weight", + "output", + "norm_weight", + "completion_counter", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_weighted.loom" + ], + "library_sources": [ + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ] + }, + { + "name": "ggml_mul_mat_id_vector_weighted_wave64_publish_q8", + "family": "loom_libs", + "symbol": "ggml_mul_mat_id_vector_weighted_wave64_publish_q8", + "source": "ops/matmul_id/vector_weighted.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_id_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "route_ids", + "route_weights", + "weight", + "output", + "norm_weight", + "completion_counter", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul_id/vector_weighted.loom" + ], + "library_sources": [ + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul_id/vector_weighted.loom", + "motifs/matmul_id/vector_q4k_q8.loom", + "motifs/matmul/q6_f32_decode.loom", + "motifs/normalization/rmsnorm_q8_1_x4.loom" + ] + }, + { + "name": "ggml_moe_build_expert_table", + "family": "loom_libs", + "symbol": "ggml_moe_build_expert_table", + "source": "ops/moe_routing_tables.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "route_ids", + "expert_table" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/moe_routing_tables.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_moe_build_expert_partition_table", + "family": "loom_libs", + "symbol": "ggml_moe_build_expert_partition_table", + "source": "ops/moe_routing_tables.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "route_count", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "bindings": [ + "expert_table", + "partition_table" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/moe_routing_tables.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_dual_q4_f32_decode", + "family": "loom_libs", + "symbol": "ggml_mul_mat_dual_q4_f32_decode", + "source": "ops/mul_mat_dual_q4_f32_decode.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "first_weight", + "second_weight", + "first_output", + "second_output" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_dual_q4_f32_decode.loom" + ], + "library_sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_f32_f32_decode_wave64", + "family": "loom_libs", + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "source": "ops/mul_mat_f32_f32_decode.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_add_f32_f32_decode_wave64", + "family": "loom_libs", + "symbol": "ggml_mul_mat_add_f32_f32_decode_wave64", + "source": "ops/mul_mat_add_f32_f32_decode.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_add_f32_f32_decode.loom" + ], + "library_sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_mul_mat_swiglu_f32_f32_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_f32_f32_wmma", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "up_weight", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", + "source": "ops/mul_mat_swiglu_f32_f32_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_f32_f32_wmma.loom" + ], + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/mul_mat_f32_f32_wmma_core.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", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_mul_mat_swiglu_f32_f32_decode_wave64", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_f32_f32_decode_wave64", + "source": "ops/mul_mat_swiglu_f32_f32_decode.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_f32_f32_decode.loom" + ], + "library_sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/elementwise/binary_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": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/elementwise/binary_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_quantize_f32_symmetric_i4_k32", + "family": "loom_libs", + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "source": "ops/mul_mat_swiglu_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_swiglu_symmetric_i4_wmma_q8_plane", + "family": "loom_libs", + "symbol": "ggml_mul_mat_swiglu_symmetric_i4_wmma_q8_plane", + "source": "ops/mul_mat_swiglu_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "gate_weight", + "up_weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_quantize_f32_symmetric_i4_k64_plane", + "family": "loom_libs", + "symbol": "ggml_quantize_f32_symmetric_i4_k64_plane", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_wmma", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "first_weight", + "second_weight", + "first_output", + "second_output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c1", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c1", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "first_weight", + "second_weight", + "first_output", + "second_output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c2", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c2", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "first_weight", + "second_weight", + "first_output", + "second_output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c3", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c3", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "first_weight", + "second_weight", + "first_output", + "second_output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c4", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c4", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "first_weight", + "second_weight", + "first_output", + "second_output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c5", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c5", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "first_weight", + "second_weight", + "first_output", + "second_output", + "quantized_values", + "scales", + "sums" + ], + "binding_access": [ + "read", + "read", + "write", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c1", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c1", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "partial", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c2", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c2", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "partial", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c3", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c3", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "partial", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c4", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c4", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "partial", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c5", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c5", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "partial", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "source": "ops/mul_mat_symmetric_i4_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "input", + "output", + "quantized_values", + "scales", + "sums", + "partial", + "completion_counters" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read", + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_quantize_f32_symmetric_i8_k256", + "family": "loom_libs", + "symbol": "ggml_quantize_f32_symmetric_i8_k256", + "source": "ops/mul_mat_symmetric_i8_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size_arg", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size_arg", + "type": "index" + } + ], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i8_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_q5_k_symmetric_i8_wmma", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q5_k_symmetric_i8_wmma", + "source": "ops/mul_mat_symmetric_i8_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_symmetric_i8_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_q5_k_q8_plane_wmmai8_token256", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q5_k_q8_plane_wmmai8_token256", + "source": "ops/mul_mat_q5_k_q8_plane_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q5_k_q8_plane_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", + "source": "ops/mul_mat_q5_k_q8_plane_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q5_k_q8_plane_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_q4_k_q8_1_x4_swiglu_wmma_token256", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q4_k_q8_1_x4_swiglu_wmma_token256", + "source": "ops/mul_mat_q5_k_q8_plane_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "q8_input", + "weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q5_k_q8_plane_wmma.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "ggml_mul_mat_q6_k_f32_wmma_prefill_wave32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q6_k_f32_wmma_prefill_wave32", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_q6_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_q6_f32_f32", + "source": "ops/matmul/vector_q6_f32_f32_kernel.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_q6_f32_f32_kernel.loom" + ], + "library_sources": [ + "motifs/matmul/vector_q6_f32_body.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_q6_f32_body.loom" + ] + }, + { + "name": "ggml_select_symmetric_i4_k32_groups", + "family": "loom_libs", + "symbol": "ggml_select_symmetric_i4_k32_groups", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "selected_groups" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "family": "loom_libs", + "symbol": "ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "weight", + "output", + "qact", + "scales", + "selected_mask" + ], + "binding_access": [ + "read", + "write", + "read", + "read", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_top_k8_f32_partitions_register", + "family": "loom_libs", + "symbol": "ggml_top_k8_f32_partitions_register", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "values", + "partial_values", + "partial_ids" + ], + "binding_access": [ + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_top_k128_f32_reduce_gather_register", + "family": "loom_libs", + "symbol": "ggml_top_k128_f32_reduce_gather_register", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "partial_values", + "partial_ids", + "selected_outputs", + "selected_values" + ], + "binding_access": [ + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_fill_negative_f32", + "family": "loom_libs", + "symbol": "ggml_fill_negative_f32", + "source": "ops/mul_mat_q6_k_packed_f16_wmma.loom", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "element_count", + "type": "index" + } + ], + "bindings": [ + "output" + ], + "binding_access": [ + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "library_sources": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/q6_k_f16.loom", + "motifs/dequant.loom", + "motifs/mul_mat_quantized_f16_prefill.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_mul_mat_vector_q6_selected_outputs_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_q6_selected_outputs_f32_f32", + "source": "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "selected_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "selected_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "selected_outputs", + "selected_values", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom" + ], + "library_sources": [ + "motifs/matmul/vector_q6_f32_body.loom", + "motifs/matmul/vector_publish_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/matmul/vector_q6_f32_body.loom", + "motifs/matmul/vector_publish_f32.loom" + ] + }, + { + "name": "ggml_flash_attention_f32_f16_wmma", + "family": "loom_libs", + "symbol": "ggml_flash_attention_f32_f16_wmma", + "source": "ops/flash_attention/prefill_f32_f16_wmma.loom", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "gate", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + { + "name": "ggml_flash_attention_f32_f16_wmma_publish_f16", + "family": "loom_libs", + "symbol": "ggml_flash_attention_f32_f16_wmma_publish_f16", + "source": "ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "gate", + "output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom" + ], + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ] + }, + { + "name": "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "family": "loom_libs", + "symbol": "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + "source": "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "workload_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "key_value_token_count", + "type": "index" + } + ], + "bindings": [ + "query", + "key", + "value", + "mask", + "partial_max", + "partial_sum", + "partial_output", + "completion_counter", + "output", + "next_q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write", + "read_write", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom" + ], + "library_sources": [ + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_flash_attention_decode_split_mask_513_of_576", + "family": "loom_libs", + "symbol": "ggml_flash_attention_decode_split_mask_513_of_576", + "source": "ops/flash_attention/decode_split_next_q8_test.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "mask" + ], + "binding_access": [ + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "flash_attention_decode_split_next_q8_test_linked", + "primary_sources": [ + "ops/flash_attention/decode_split_next_q8_test.loom" + ], + "library_sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_flash_attention_decode_split_pack_completed_q8_test", + "family": "loom_libs", + "symbol": "ggml_flash_attention_decode_split_pack_completed_q8_test", + "source": "ops/flash_attention/decode_split_next_q8_test.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "flash_attention_decode_split_next_q8_test_linked", + "primary_sources": [ + "ops/flash_attention/decode_split_next_q8_test.loom" + ], + "library_sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_flash_attention_decode_split_quantize_reference_4096", + "family": "loom_libs", + "symbol": "ggml_flash_attention_decode_split_quantize_reference_4096", + "source": "ops/flash_attention/decode_split_next_q8_test.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "input", + "output" + ], + "binding_access": [ + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "archive", + "link_module": "flash_attention_decode_split_next_q8_test_linked", + "primary_sources": [ + "ops/flash_attention/decode_split_next_q8_test.loom" + ], + "library_sources": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "ggml_get_rows_f32", + "family": "loom_libs", + "symbol": "ggml_get_rows_f32", + "source": "ops/rows/get_rows_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "library_sources": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_get_rows_scale_f32", + "family": "loom_libs", + "symbol": "ggml_get_rows_scale_f32", + "source": "ops/rows/get_rows_scale_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "raw_output", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/get_rows_scale_f32.loom" + ], + "library_sources": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_gather_add_rmsnorm_binary_f32", + "family": "loom_libs", + "symbol": "ggml_gather_add_rmsnorm_binary_f32", + "source": "ops/rows/gather_add_rmsnorm_binary_f32.loom", + "workload_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "source_token_count", + "type": "index" + }, + { + "name": "output_token_count", + "type": "index" + } + ], + "bindings": [ + "attention", + "residual", + "output_ids", + "raw_output", + "weight", + "normalized_output" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/gather_add_rmsnorm_binary_f32.loom" + ], + "library_sources": [ + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "family": "loom_libs", + "symbol": "ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "source": "ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "norm_weight", + "raw_output", + "output", + "q8_output", + "f16_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "read_write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom" + ], + "library_sources": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/normalization/publish_q8.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/normalization/publish_q8.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_get_rows_f32_next", + "family": "loom_libs", + "symbol": "ggml_get_rows_f32_next", + "source": "ops/rows/get_rows_next_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "output", + "next_output" + ], + "binding_access": [ + "read", + "read", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/get_rows_next_f32.loom" + ], + "library_sources": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + "compile_dependencies": [ + "motifs/rows/get_rows_decode.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ] + }, + { + "name": "ggml_rope_f32", + "family": "loom_libs", + "symbol": "ggml_rope_f32", + "source": "ops/rope/rope_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_span", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_span", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_rope_scale_f32", + "family": "loom_libs", + "symbol": "ggml_rope_scale_f32", + "source": "ops/rope/rope_scale_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_scale_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_rope_token_scale_f32", + "family": "loom_libs", + "symbol": "ggml_rope_token_scale_f32", + "source": "ops/rope/rope_scale_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "scales", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_scale_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_rope_concat_f32", + "family": "loom_libs", + "symbol": "ggml_rope_concat_f32", + "source": "ops/rope/rope_concat_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "input", + "theta", + "freq_factors", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_concat_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom" + ] + }, + { + "name": "ggml_set_rows", + "family": "loom_libs", + "symbol": "ggml_set_rows", + "source": "ops/rows/set_rows.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "rows", + "indices", + "cache" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "library_sources": [ + "motifs/rows/cache_publish.loom" + ] + }, + "compile_dependencies": [ + "motifs/rows/cache_publish.loom" + ] + }, + { + "name": "ggml_rope_set_rows_f32", + "family": "loom_libs", + "symbol": "ggml_rope_set_rows_f32", + "source": "ops/rope/rope_set_rows_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "indices", + "input", + "theta", + "freq_factors", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/rope/rope_set_rows_f32.loom" + ], + "library_sources": [ + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" + ] + }, + "compile_dependencies": [ + "motifs/rope/rope_f32.loom", + "motifs/rows/cache_publish.loom" + ] + }, + { + "name": "llm_ssm_conv_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_f32", + "source": "ops/ssm_conv/generic_f32.loom", + "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", + "output" + ], + "binding_access": [ + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/generic_f32.loom" + ], + "library_sources": [ + "motifs/ssm_conv/generic_apply.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/generic_apply.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "llm_ssm_conv_lfm_dconv3_state_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_lfm_dconv3_state_f32", + "source": "ops/ssm_conv/lfm_dconv3_state_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state", + "x", + "filter", + "output", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/lfm_dconv3_state_f32.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [] + }, + { + "name": "llm_ssm_conv_binary_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_binary_f32", + "source": "ops/ssm_conv/generic_binary_f32.loom", + "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", + "operand", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/generic_binary_f32.loom" + ], + "library_sources": [ + "motifs/ssm_conv/generic_apply.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/generic_apply.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/elementwise/binary_apply.loom" + ] + }, + { + "name": "llm_ssm_conv_snapshot_window_tail_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_snapshot_window_tail_f32", + "source": "ops/ssm_conv/snapshot_window_tail_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state", + "x", + "snapshot", + "cache" + ], + "binding_access": [ + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/snapshot_window_tail_f32.loom" + ], + "library_sources": [ + "motifs/ssm_conv/snapshot_window_tail.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/snapshot_window_tail.loom" + ] + }, + { + "name": "llm_ssm_conv_dconv4_silu_prefill_finish_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_dconv4_silu_prefill_finish_f32", + "source": "ops/ssm_conv/prefill_finish_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state", + "filter", + "edges", + "output", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/prefill_finish_f32.loom" + ], + "library_sources": [ + "motifs/ssm_conv/prefill_finish.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/prefill_finish.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ] + }, + { + "name": "llm_ssm_conv_dconv4_silu_prefill_512_wg1024", + "family": "loom_libs", + "symbol": "llm_ssm_conv_dconv4_silu_prefill_512_wg1024", + "source": "ops/ssm_conv/prefill_dconv4_silu_512_wg1024.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state_snapshot", + "x", + "filter", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/prefill_dconv4_silu_512_wg1024.loom" + ], + "library_sources": [ + "motifs/ssm_conv/prefill_window.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/prefill_window.loom" + ] + }, + { + "name": "llm_ssm_conv_dconv4_silu_decode_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_dconv4_silu_decode_f32", + "source": "ops/ssm_conv/decode_dconv4_silu_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state", + "x", + "filter", + "output", + "cache" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/decode_dconv4_silu_f32.loom" + ], + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ] + }, + { + "name": "llm_ssm_conv_dconv4_silu_rollback_f32", + "family": "loom_libs", + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "source": "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "state", + "x", + "filter", + "output", + "cache0", + "cache1", + "cache2", + "cache3", + "cache4" + ], + "binding_access": [ + "read", + "read", + "read", + "write", + "write", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ] + }, + "compile_dependencies": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128", + "source": "ops/gated_delta_net/head128_scan_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "g", + "beta", + "state_in", + "dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_scan_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_rmsnorm_gate", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_rmsnorm_gate", + "source": "ops/gated_delta_net/head128_scan_rmsnorm_gate_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "g", + "beta", + "state_in", + "dst", + "rms_weight", + "raw_gate", + "norm_output", + "half_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "write", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_scan_rmsnorm_gate_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_projection_epilogue", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_projection_epilogue", + "source": "ops/gated_delta_net/head128_scan_projection_epilogue_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "alpha_raw", + "beta_raw", + "bias", + "a_scale", + "state_in", + "dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_scan_projection_epilogue_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_inplace", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_inplace", + "source": "ops/gated_delta_net/head128_inplace_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "g", + "beta", + "state_inout", + "dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_inplace_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_inplace_projection_epilogue", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_inplace_projection_epilogue", + "source": "ops/gated_delta_net/head128_inplace_projection_epilogue_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "alpha_raw", + "beta_raw", + "bias", + "a_scale", + "state_inout", + "dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read_write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_inplace_projection_epilogue_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "source": "ops/gated_delta_net/head128_snapshot_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "g", + "beta", + "state_in", + "snapshot_cache", + "dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_snapshot_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue", + "source": "ops/gated_delta_net/head128_snapshot_projection_epilogue_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "alpha_raw", + "beta_raw", + "bias", + "a_scale", + "state_in", + "snapshot_cache", + "dst" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_snapshot_projection_epilogue_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_epilogue", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_epilogue", + "source": "ops/gated_delta_net/head128_selected_snapshot_projection_epilogue_f32_wmma.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "alpha_raw", + "beta_raw", + "bias", + "a_scale", + "state_in", + "snapshot_cache", + "dst", + "state_ids" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "write", + "write", + "read" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_selected_snapshot_projection_epilogue_f32_wmma.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_rms_gate_q8", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_rms_gate_q8", + "source": "ops/gated_delta_net/head128_selected_snapshot_projection_rms_gate_publish_q8.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "q", + "k", + "v", + "alpha_raw", + "beta_raw", + "bias", + "a_scale", + "state_in", + "snapshot_cache", + "dst", + "state_ids", + "rms_weight", + "raw_gate", + "q8_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "read", + "write", + "write", + "read", + "read", + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/gated_delta_net/head128_selected_snapshot_projection_rms_gate_publish_q8.loom" + ], + "library_sources": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + "compile_dependencies": [ + "motifs/gated_delta_net/scan_state_head128.loom", + "motifs/gated_delta_net/projection_epilogue.loom", + "motifs/gated_delta_net/rmsnorm_gate.loom", + "motifs/gated_delta_net/inplace_state_head128.loom", + "motifs/gated_delta_net/snapshot_load.loom", + "motifs/gated_delta_net/snapshot_publish.loom", + "motifs/gated_delta_net/selected_snapshot_load.loom", + "motifs/gated_delta_net/publish_f32.loom", + "motifs/gated_delta_net/publish_q8.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/elementwise/unary_apply.loom", + "motifs/quantize_q8_1_x4.loom" + ] + }, + { + "name": "llm_gated_delta_net_projection_epilogue_f32", + "family": "loom_libs", + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "source": "ops/gated_delta_net/projection_epilogue_f32.loom", + "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": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/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_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_aligned", + "family": "loom_libs", + "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", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/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_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_bias_residual_publish_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", + "source": "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/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_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_bias_residual_publish_f32_aligned", + "family": "loom_libs", + "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", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "gate_weight", + "up_weight", + "bias", + "residual_input", + "residual_output" + ], + "binding_access": [ + "read", + "read", + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom" + ], + "library_sources": [ + "motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom", + "motifs/matmul/tiled_pair_postprocess_f32.loom", + "motifs/matmul/tiled_postprocess_f32.loom", + "motifs/matmul/tiled_publish.loom", + "motifs/dequant.loom", + "motifs/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_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_vector_iq3_xxs_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_iq3_xxs_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", + "weight", + "iq3_xxs_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_iq3_xxs_publish_f32", + "family": "loom_libs", + "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", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "iq3_xxs_grid", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" + ], + "library_sources": [ + "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/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_mul_mat_vector_iq1_s_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_iq1_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", + "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_vector_iq1_m_f32_f32", + "family": "loom_libs", + "symbol": "ggml_mul_mat_vector_iq1_m_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", + "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": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" + ], + "library_sources": [ + "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/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_mul_mat_tiled_input_f32_iq1_m_publish_f32", + "family": "loom_libs", + "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", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + } + ], + "bindings": [ + "input", + "weight", + "iq1_grid", + "output" + ], + "binding_access": [ + "read", + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "ops/matmul/tiled_input_f32_codebook_publish_f32.loom" + ], + "library_sources": [ + "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/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" + ] + } + ], + "link_modules": [ + { + "name": "flash_attention_decode_split_next_q8_test_linked", + "srcs": [ + "ops/flash_attention/decode_split_next_q8_test.loom" + ], + "libraries": [ + "ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom", + "motifs/flash_attention/decode_partials.loom", + "motifs/flash_attention/decode_reduce.loom", + "motifs/flash_attention/publish_f32.loom", + "motifs/flash_attention/publish_q8.loom", + "motifs/flash_attention/completion_counter_reduce.loom", + "motifs/quantize_q8_1_x4.loom" + ] + } + ], + "plan_cases": [ + { + "name": "get_rows_f32_q4k_decode_plan_test", + "args": [ + "$(location ops/rows/get_rows_f32.loom)", + "--benchmark=@ggml_get_rows_f32_q4k_decode", + "--config=ggml.get_rows_f32.token_capacity=1", + "--config=ggml.get_rows_f32.hidden_capacity=2048", + "--config=ggml.get_rows_f32.weight_format=4", + "--dry-run", + "--output-format=jsonl" + ], + "source": "ops/rows/get_rows_f32.loom" + }, + { + "name": "get_rows_f32_q4k_row_access_plan_test", + "args": [ + "$(location ops/rows/get_rows_f32.loom)", + "--benchmark=@ggml_get_rows_f32_q4k_row_access", + "--config=ggml.get_rows_f32.token_capacity=3", + "--config=ggml.get_rows_f32.hidden_capacity=2048", + "--config=ggml.get_rows_f32.weight_format=4", + "--dry-run", + "--output-format=jsonl" + ], + "source": "ops/rows/get_rows_f32.loom" + }, + { + "name": "get_rows_scale_f32_plan_test", + "args": [ + "$(location ops/rows/get_rows_scale_f32.loom)", + "--benchmark=@ggml_get_rows_scale_f32_f32", + "--config=ggml.get_rows_f32.token_capacity=4", + "--config=ggml.get_rows_f32.hidden_capacity=4", + "--config=ggml.get_rows_f32.weight_format=32", + "--config=ggml.get_rows_scale_f32.scale=2.0", + "--config=ggml.get_rows_scale_f32.bias=0.25", + "--dry-run", + "--output-format=jsonl" + ], + "source": "ops/rows/get_rows_scale_f32.loom" + } + ], + "sources": [], + "kernels": [] +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom new file mode 100644 index 000000000000..876b2be71732 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom @@ -0,0 +1,168 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.attention.postprocess.head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) + +config.decl @qwen3_moe.model.rms_epsilon : f32 + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.rope_mscale : f32 + +func.def inline @qwen3_moe_rmsnorm_neox_packet(%row_sum: f32, %head_size: f32, %epsilon: f32, %position: f32, %rope_mscale: f32, %inverse_frequencies: vector<2xf32>, %low_values: vector<2xf32>, %high_values: vector<2xf32>, %low_weights: vector<2xf32>, %high_weights: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %mean = scalar.divf %row_sum, %head_size : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + %scale_vector = vector.splat %scale : vector<2xf32> + %position_vector = vector.splat %position : vector<2xf32> + %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 + %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> + %normalized_low = vector.mulf %low_values, %scale_vector : vector<2xf32> + %normalized_high = vector.mulf %high_values, %scale_vector : vector<2xf32> + %scaled_low = vector.mulf %normalized_low, %low_weights : vector<2xf32> + %scaled_high = vector.mulf %normalized_high, %high_weights : vector<2xf32> + %angles = vector.mulf %position_vector, %inverse_frequencies : vector<2xf32> + %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> + %cosines = vector.costurnsf %turns : vector<2xf32> + %sines = vector.sinturnsf %turns : vector<2xf32> + %low_cosines = vector.mulf %scaled_low, %cosines : vector<2xf32> + %high_sines = vector.mulf %scaled_high, %sines : vector<2xf32> + %low_sines = vector.mulf %scaled_low, %sines : vector<2xf32> + %high_cosines = vector.mulf %scaled_high, %cosines : vector<2xf32> + %rotated_low0 = vector.subf %low_cosines, %high_sines : vector<2xf32> + %rotated_high0 = vector.addf %low_sines, %high_cosines : vector<2xf32> + %rope_mscale_vector = vector.splat %rope_mscale : vector<2xf32> + %rotated_low = vector.mulf %rotated_low0, %rope_mscale_vector : vector<2xf32> + %rotated_high = vector.mulf %rotated_high0, %rope_mscale_vector : vector<2xf32> + func.return %rotated_low, %rotated_high : vector<2xf32>, vector<2xf32> +} +// Processes one head domain after its raw projection row is visible. Callers +// may use a larger workgroup than the packet count; inactive workitems +// contribute zero to normalization and perform no loads or stores. +template.def<@qwen3_moe.attention.postprocess.head_body> device @qwen3_moe_attention_postprocess_head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) { + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %rope_mscale = config.get @qwen3_moe.attention.rope_mscale : f32 + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), range(%head_size0, 4, 1024), mul(%head_size0, 4), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %half_head_size = index.div %head_size, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain = index.assume %head_domain0 [lt(%head_domain0, %head_domain_count)] : index + %key_domain_end = index.add %query_head_count, %key_value_head_count : index + %is_query = index.cmp ult, %head_domain, %query_head_count : index + %is_query_or_key = index.cmp ult, %head_domain, %key_domain_end : index + %key_value_head = index.rem %head_domain, %key_value_head_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %head_size_i32 = index.cast %head_size : index to i32 + %head_size_f32 = scalar.sitofp %head_size_i32 : i32 to f32 + %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %query_input_noalias, %key_input_noalias, %value_input_noalias, %query_norm_weight_noalias, %key_norm_weight_noalias, %inverse_frequencies_noalias, %query_output_noalias, %key_cache_noalias, %value_cache_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %query_input_view = buffer.view %query_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_input_view = buffer.view %key_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %value_input_view = buffer.view %value_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> + %query_norm_weight_view = buffer.view %query_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %key_norm_weight_view = buffer.view %key_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> + %inverse_frequencies_view = buffer.view %inverse_frequencies_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + %key_cache_words_view = buffer.view %key_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %value_cache_words_view = buffer.view %value_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + %row_sum = scf.if %is_query_or_key -> (f32) { + %partial_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %active_channel -> (f32) { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %reduction_channel = index.mul %channel, %c4 : index + %reduction_values = scf.if %is_query -> (vector<4xf32>) { + %query_values = vector.load %query_input_view[%token, %head_domain, %reduction_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %query_values : vector<4xf32> + } else { + %key_values = vector.load %key_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + scf.yield %key_values : vector<4xf32> + } + %squares = vector.mulf %reduction_values, %reduction_values : vector<4xf32> + %sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %sum = kernel.workgroup.reduce %partial_sum : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.if %publish_output { + scf.if %active_channel { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %paired_channel = index.add %pair_channel, %half_head_size : index + %paired_word = index.add %channel, %pair_packet_count : index + %reduction_channel = index.mul %channel, %c4 : index + scf.if %is_query { + %low_values = vector.load %query_input_view[%token, %head_domain, %pair_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %query_input_view[%token, %head_domain, %paired_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %query_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %query_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %rope_mscale, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + vector.store %rotated_low, %query_output_view[%token, %head_domain, %pair_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + vector.store %rotated_high, %query_output_view[%token, %head_domain, %paired_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> + } else { + scf.if %is_query_or_key { + %low_values = vector.load %key_input_view[%token, %key_value_head, %pair_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %high_values = vector.load %key_input_view[%token, %key_value_head, %paired_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> + %low_weights = vector.load %key_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> + %high_weights = vector.load %key_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> + %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %cache_index_raw = view.load %key_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %rope_mscale, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %half_low = vector.fptrunc %rotated_low : vector<2xf32> to vector<2xf16> + %half_high = vector.fptrunc %rotated_high : vector<2xf32> to vector<2xf16> + %packed_low = vector.bitcast %half_low : vector<2xf16> to vector<1xi32> + %packed_high = vector.bitcast %half_high : vector<2xf16> to vector<1xi32> + vector.store %packed_low, %key_cache_words_view[%cache_index, %key_value_head, %channel] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + vector.store %packed_high, %key_cache_words_view[%cache_index, %key_value_head, %paired_word] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } else { + %cache_index_raw = view.load %value_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 + %cache_index0 = index.cast %cache_index_i64 : i64 to index + %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index + %values = vector.load %value_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> + %half_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + %packed_values = vector.bitcast %half_values : vector<4xf16> to vector<2xi32> + vector.store %packed_values, %value_cache_words_view[%cache_index, %key_value_head, %pair_channel] : vector<2xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> + } + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/bf16_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/bf16_f16.loom new file mode 100644 index 000000000000..1a38444c8476 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/bf16_f16.loom @@ -0,0 +1,12 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Loads four adjacent BF16 weights and truncates them for FP16 matrix staging. +func.def inline @ggml_bf16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xbf16> + %values_bf16 = vector.load %row_view[%k] : view<[%bounded_input_size]xbf16> -> vector<4xbf16> + %values_f32 = vector.extf %values_bf16 : vector<4xbf16> to vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} 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 new file mode 100644 index 000000000000..c9526fb51851 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/concat_dim0.loom @@ -0,0 +1,149 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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_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)] +config.decl @ggml.concat_dim0_strided_source_f32.row_count : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.concat_dim0_strided_source_f32.row_ne1 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.concat_dim0_strided_source_f32.row_ne2 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride1 : %value: index where [range(%value, 1, 1073741824)] +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(%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(%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 + %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 + %element_count = index.assume %element_count0 [range(%element_count0, 1, 1073741824)] : index + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workgroup = index.assume %workgroup0 [range(%workgroup0, 0, 4194303)] : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %linear_base = index.mul %workgroup, %workgroup_size : index + %linear0 = index.add %linear_base, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 1073742079)] : index + %in_bounds = index.cmp ult, %linear, %element_count : index + + %lhs_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %lhs, %rhs, %output : buffer, buffer, buffer + %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%element_count]xf32> + + scf.if %in_bounds { + %row = index.div %linear, %output_width : index + %column = index.rem %linear, %output_width : index + %from_lhs = index.cmp ult, %column, %lhs_width : index + %value = scf.if %from_lhs -> (f32) { + %lhs_row = index.mul %row, %lhs_width : index + %lhs_index0 = index.add %lhs_row, %column : index + %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, %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<[%element_count]xf32> + } + template.return +} + +template.def<@ggml.copy.concat_dim0.strided_source_f32_body> device @ggml_copy_concat_dim0_strided_source_f32_body(%lhs_span: index, %rhs_span: 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_strided_source_f32.lhs_width : index + %rhs_width0 = config.get @ggml.concat_dim0_strided_source_f32.rhs_width : index + %row_count0 = config.get @ggml.concat_dim0_strided_source_f32.row_count : index + %row_ne1 = config.get @ggml.concat_dim0_strided_source_f32.row_ne1 : index + %row_ne2 = config.get @ggml.concat_dim0_strided_source_f32.row_ne2 : index + %lhs_s0 = config.get @ggml.concat_dim0_strided_source_f32.lhs_stride0 : index + %lhs_s1 = config.get @ggml.concat_dim0_strided_source_f32.lhs_stride1 : index + %lhs_s2 = config.get @ggml.concat_dim0_strided_source_f32.lhs_stride2 : index + %lhs_s3 = config.get @ggml.concat_dim0_strided_source_f32.lhs_stride3 : index + %rhs_s0 = config.get @ggml.concat_dim0_strided_source_f32.rhs_stride0 : index + %rhs_s1 = config.get @ggml.concat_dim0_strided_source_f32.rhs_stride1 : index + %rhs_s2 = config.get @ggml.concat_dim0_strided_source_f32.rhs_stride2 : index + %rhs_s3 = config.get @ggml.concat_dim0_strided_source_f32.rhs_stride3 : 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 + %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 + %element_count = index.assume %element_count0 [range(%element_count0, 1, 1073741824)] : index + %lhs_span_checked = index.assume %lhs_span [range(%lhs_span, 1, 1073741824)] : index + %rhs_span_checked = index.assume %rhs_span [range(%rhs_span, 1, 1073741824)] : index + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workgroup = index.assume %workgroup0 [range(%workgroup0, 0, 4194303)] : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %linear_base = index.mul %workgroup, %workgroup_size : index + %linear0 = index.add %linear_base, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 1073742079)] : index + %in_bounds = index.cmp ult, %linear, %element_count : index + + %lhs_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %lhs, %rhs, %output : buffer, buffer, buffer + %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%element_count]xf32> + + scf.if %in_bounds { + %row = index.div %linear, %output_width : index + %column = index.rem %linear, %output_width : index + %coord1 = index.rem %row, %row_ne1 : index + %row_after_ne1 = index.div %row, %row_ne1 : index + %coord2 = index.rem %row_after_ne1, %row_ne2 : index + %coord3 = index.div %row_after_ne1, %row_ne2 : index + %from_lhs = index.cmp ult, %column, %lhs_width : index + %value = scf.if %from_lhs -> (f32) { + %offset0 = index.mul %column, %lhs_s0 : index + %offset1 = index.mul %coord1, %lhs_s1 : index + %offset2 = index.mul %coord2, %lhs_s2 : index + %offset3 = index.mul %coord3, %lhs_s3 : index + %offset01 = index.add %offset0, %offset1 : index + %offset23 = index.add %offset2, %offset3 : index + %lhs_index0 = index.add %offset01, %offset23 : index + %lhs_index, %lhs_view_span = index.assume %lhs_index0, %lhs_span_checked [lt(%lhs_index0, %lhs_span_checked)] : index, index + %lhs_view = buffer.view %lhs_noalias[%base] : buffer -> view<[%lhs_view_span]xf32> + %loaded = view.load %lhs_view[%lhs_index] : view<[%lhs_view_span]xf32> -> f32 + scf.yield %loaded : f32 + } else { + %rhs_column = index.sub %column, %lhs_width : index + %offset0 = index.mul %rhs_column, %rhs_s0 : index + %offset1 = index.mul %coord1, %rhs_s1 : index + %offset2 = index.mul %coord2, %rhs_s2 : index + %offset3 = index.mul %coord3, %rhs_s3 : index + %offset01 = index.add %offset0, %offset1 : index + %offset23 = index.add %offset2, %offset3 : index + %rhs_index0 = index.add %offset01, %offset23 : index + %rhs_index, %rhs_view_span = index.assume %rhs_index0, %rhs_span_checked [lt(%rhs_index0, %rhs_span_checked)] : index, index + %rhs_view = buffer.view %rhs_noalias[%base] : buffer -> view<[%rhs_view_span]xf32> + %loaded = view.load %rhs_view[%rhs_index] : view<[%rhs_view_span]xf32> -> f32 + scf.yield %loaded : f32 + } + 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/convert_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/convert_f32_f16.loom new file mode 100644 index 000000000000..a99148132e5d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/convert_f32_f16.loom @@ -0,0 +1,33 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.convert_f32_f16.body(%element_count: index, %source: buffer, %output: buffer) +template.decl @ggml.copy.convert_f32_f16.vector8(%values: vector<8xf32>) -> (vector<8xf16>) + +template.def<@ggml.copy.convert_f32_f16.body> device @ggml_copy_convert_f32_f16_body(%element_count: index, %source: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 4, 1073741824), mul(%element_count, 4)] : index + %group = kernel.workgroup.id : index + %thread = kernel.workitem.id : index + %threads = index.constant 256 : index + %four = index.constant 4 : index + %lane = index.madd %group, %threads, %thread : index + %linear0 = index.mul %lane, %four : index + %linear = index.assume %linear0 [range(%linear0, 0, 1073742844), mul(%linear0, 4)] : index + %end = index.add %linear, %four : index + %valid = index.cmp ule, %end, %count : index + %zero = index.constant 0 : offset + %source_noalias, %output_noalias = buffer.assume.noalias %source, %output : buffer, buffer + %source_view = buffer.view %source_noalias[%zero] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output_noalias[%zero] : buffer -> view<[%count]xf16> + scf.if %valid { + %values = vector.load %source_view[%linear] : view<[%count]xf32> -> vector<4xf32> + %converted = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store %converted, %output_view[%linear] : vector<4xf16>, view<[%count]xf16> + } + template.return +} + +template.def<@ggml.copy.convert_f32_f16.vector8> device @ggml_copy_convert_f32_f16_vector8(%values: vector<8xf32>) -> (vector<8xf16>) { + %converted = vector.fptrunc %values : vector<8xf32> to vector<8xf16> + template.return %converted : vector<8xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/layout_k16_major.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/layout_k16_major.loom new file mode 100644 index 000000000000..16d4b6c84544 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/layout_k16_major.loom @@ -0,0 +1,60 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +config.decl @ggml.copy_f16_k16_major.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] +config.decl @ggml.copy_f16_k16_major.token_count : %value: index where [range(%value, 512, 2048), mul(%value, 512)] +config.decl @ggml.copy_f16_k16_major.input_is_f16 : %value: index where [range(%value, 0, 1)] + +template.decl @ggml.copy.convert_f32_f16.vector8(%values: vector<8xf32>) -> (vector<8xf16>) +template.decl @ggml.copy.layout_k16_major.body(%token_count: index, %input: buffer, %output: buffer) + +template.def<@ggml.copy.layout_k16_major.body> device @ggml_copy_layout_k16_major_body(%token_count: index, %input: buffer, %output: buffer) { + %width = config.get @ggml.copy_f16_k16_major.input_size : index + %token_count_static = config.get @ggml.copy_f16_k16_major.token_count : index + %tokens = index.assume %token_count_static [range(%token_count_static, 16, 2048), mul(%token_count_static, 16)] : index + %count = index.mul %tokens, %width : index + %c0 = index.constant 0 : offset + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %thread = kernel.workitem.id : index + %group = kernel.workgroup.id : index + %packet = index.madd %group, %c256, %thread : index + %linear_raw = index.mul %packet, %c8 : index + %linear = index.assume %linear_raw [range(%linear_raw, 0, 67110904), mul(%linear_raw, 8)] : index + %end = index.add %linear, %c8 : index + %valid = index.cmp ule, %end, %count : index + scf.if %valid { + %tile = index.div %linear, %c256 : index + %local = index.rem %linear, %c256 : index + %m_tiles_raw = index.div %tokens, %c16 : index + %m_tiles = index.assume %m_tiles_raw [range(%m_tiles_raw, 1, 128)] : index + %m_tile = index.rem %tile, %m_tiles : index + %k_tile = index.div %tile, %m_tiles : index + %row_in_tile = index.div %local, %c16 : index + %k_in_tile_raw = index.rem %local, %c16 : index + %k_in_tile = index.assume %k_in_tile_raw [range(%k_in_tile_raw, 0, 8), mul(%k_in_tile_raw, 8)] : index + %row = index.madd %m_tile, %c16, %row_in_tile : index + %k = index.madd %k_tile, %c16, %k_in_tile : index + %bounded_row, %bounded_tokens = index.assume %row, %tokens [lt(%row, %tokens)] : index, index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %dest = buffer.view %output_noalias[%c0] : buffer -> view<[%count]xf16> + %maximum_k = index.sub %width, %c8 : index + %bounded_k = index.assume %k [range(%k, 0, %maximum_k), mul(%k, 8)] : index + %input_is_f16 = config.get @ggml.copy_f16_k16_major.input_is_f16 : index + %one = index.constant 1 : index + %use_f16 = index.cmp eq, %input_is_f16, %one : index + %values = scf.if %use_f16 -> (vector<8xf16>) { + %source = buffer.view %input_noalias[%c0] : buffer -> view<[%tokens]x[%width]xf16> + %loaded = vector.load %source[%bounded_row, %bounded_k] : view<[%tokens]x[%width]xf16> -> vector<8xf16> + scf.yield %loaded : vector<8xf16> + } else { + %source = buffer.view %input_noalias[%c0] : buffer -> view<[%tokens]x[%width]xf32> + %loaded = vector.load %source[%bounded_row, %bounded_k] : view<[%tokens]x[%width]xf32> -> vector<8xf32> + %converted = template.apply<@ggml.copy.convert_f32_f16.vector8>(%loaded) : (vector<8xf32>) -> (vector<8xf16>) + scf.yield %converted : vector<8xf16> + } + vector.store %values, %dest[%linear] : vector<8xf16>, view<[%count]xf16> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/strided_source.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/strided_source.loom new file mode 100644 index 000000000000..d656390eb83c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/strided_source.loom @@ -0,0 +1,54 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +config.decl @ggml.copy_strided_source_f32.ne0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.ne1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride3 : %value: index where [range(%value, 1, 1073741824)] + +template.decl @ggml.copy.strided_source.f32_body(%element_count: index, %source_span: index, %source: buffer, %output: buffer) + +template.def<@ggml.copy.strided_source.f32_body> device @ggml_copy_strided_source_f32_body(%element_count: index, %source_span: index, %source: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 1073741824)] : index + %n0 = config.get @ggml.copy_strided_source_f32.ne0 : index + %n1 = config.get @ggml.copy_strided_source_f32.ne1 : index + %n2 = config.get @ggml.copy_strided_source_f32.ne2 : index + %span = index.assume %source_span [range(%source_span, 1, 1073741824)] : index + %s0 = config.get @ggml.copy_strided_source_f32.stride0 : index + %s1 = config.get @ggml.copy_strided_source_f32.stride1 : index + %s2 = config.get @ggml.copy_strided_source_f32.stride2 : index + %s3 = config.get @ggml.copy_strided_source_f32.stride3 : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 1073742079)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %source_noalias, %output_noalias = buffer.assume.noalias %source, %output : buffer, buffer + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %coord0 = index.rem %linear, %n0 : index + %remaining0 = index.div %linear, %n0 : index + %coord1 = index.rem %remaining0, %n1 : index + %remaining1 = index.div %remaining0, %n1 : index + %coord2 = index.rem %remaining1, %n2 : index + %coord3 = index.div %remaining1, %n2 : index + %offset0 = index.mul %coord0, %s0 : index + %offset1 = index.mul %coord1, %s1 : index + %offset2 = index.mul %coord2, %s2 : index + %offset3 = index.mul %coord3, %s3 : index + %offset01 = index.add %offset0, %offset1 : index + %offset23 = index.add %offset2, %offset3 : index + %source_index0 = index.add %offset01, %offset23 : index + %source_index, %view_span = index.assume %source_index0, %span [lt(%source_index0, %span)] : index, index + %source_view = buffer.view %source_noalias[%zero_offset] : buffer -> view<[%view_span]xf32> + %value = view.load %source_view[%source_index] : view<[%view_span]xf32> -> f32 + view.store %value, %output_view[%linear] : f32, view<[%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 new file mode 100644 index 000000000000..bac75d38b078 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/transpose_tile.loom @@ -0,0 +1,56 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.transpose_tile.f16_body(%rows: index, %columns: index, %input: buffer, %output: buffer) + +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 + %c32 = index.constant 32 : index + %zero = index.constant 0 : offset + %bytes = index.constant 2560 : offset + %wi = kernel.workitem.id : index + %gx = kernel.workgroup.id : index + %gy = kernel.workgroup.id : index + %row_origin = index.mul %gy, %c32 : index + %col_origin = index.mul %gx, %c32 : index + %row = index.div %wi, %c8 : index + %col_word = index.rem %wi, %c8 : index + %col = index.mul %col_word, %c4 : 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<[%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<[%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 + %read_col0 = index.add %col, %add0 : index + %s0 = view.load %shared_transpose_view[%row, %read_col0] : view<32x32xf16, %transposed> -> f16 + %add1 = index.constant 1 : index + %read_col1 = index.add %col, %add1 : index + %s1 = view.load %shared_transpose_view[%row, %read_col1] : view<32x32xf16, %transposed> -> f16 + %add2 = index.constant 2 : index + %read_col2 = index.add %col, %add2 : index + %s2 = view.load %shared_transpose_view[%row, %read_col2] : view<32x32xf16, %transposed> -> f16 + %add3 = index.constant 3 : index + %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_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 + %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/copy/vector_load_store.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/vector_load_store.loom new file mode 100644 index 000000000000..24ae28d7fbb3 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/copy/vector_load_store.loom @@ -0,0 +1,24 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.vector_load_store.f32_body(%element_count: index, %source: buffer, %output: buffer) + +template.def<@ggml.copy.vector_load_store.f32_body> device @ggml_copy_vector_load_store_f32_body(%element_count: index, %source: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 1073741824)] : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 1073742079)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %source_noalias, %output_noalias = buffer.assume.noalias %source, %output : buffer, buffer + %source_view = buffer.view %source_noalias[%zero_offset] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %value = view.load %source_view[%linear] : view<[%count]xf32> -> f32 + view.store %value, %output_view[%linear] : f32, view<[%count]xf32> + } + 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 new file mode 100644 index 000000000000..b7988bb3babb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom @@ -0,0 +1,3412 @@ +func.def inline @ggml_paired_weight_words(%half_packet: i1, %paired: i1, %load_up: i1, %weight: buffer, %peer: buffer, %offset: offset) -> (vector<4xi32>) { + %zero = index.constant 0 : index + %padding = vector.constant 0 : vector<2xi32> + %weight_words = scf.if %half_packet -> (vector<4xi32>) { + %view = buffer.view %weight[%offset] : buffer -> view<2xi32> + %loaded = vector.load %view[%zero] : view<2xi32> -> vector<2xi32> + %words = vector.concat<0> %loaded, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } else { + %view = buffer.view %weight[%offset] : buffer -> view<4xi32> + %words = vector.load %view[%zero] : view<4xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %peer_words = scf.if %paired -> (vector<4xi32>) { + %words = scf.if %half_packet -> (vector<4xi32>) { + %view = buffer.view %peer[%offset] : buffer -> view<2xi32> + %loaded = vector.load %view[%zero] : view<2xi32> -> vector<2xi32> + %packet = vector.concat<0> %loaded, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %packet : vector<4xi32> + } else { + %view = buffer.view %peer[%offset] : buffer -> view<4xi32> + %packet = vector.load %view[%zero] : view<4xi32> -> vector<4xi32> + scf.yield %packet : vector<4xi32> + } + scf.yield %words : vector<4xi32> + } else { + scf.yield %weight_words : vector<4xi32> + } + %words = scf.select %load_up, %peer_words, %weight_words : vector<4xi32> + func.return %words : vector<4xi32> +} + +func.def public inline @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %field_bytes = index.constant 1024 : offset + %false = scalar.constant false : i1 + %header = func.call @ggml_paired_weight_words(%false, %paired, %is_up, %weight, %peer, %row_base) : (i1, i1, i1, buffer, buffer, offset) -> (vector<4xi32>) + %halves = vector.bitcast %header : vector<4xi32> to vector<8xf16> + %d_half = vector.extract %halves[0] : vector<8xf16> -> f16 + %m_half = vector.extract %halves[1] : vector<8xf16> -> f16 + %d = scalar.extf %d_half : f16 to f32 + %m = scalar.extf %m_half : f16 to f32 + %h0 = vector.extract %header[1] : vector<4xi32> -> i32 + %h1 = vector.extract %header[2] : vector<4xi32> -> i32 + %h2 = vector.extract %header[3] : vector<4xi32> -> i32 + %group1 = index.add %group0, %c1 : index + %s0, %m0 = func.call @ggml_q4k_scale_min_from_header(%h0, %h1, %h2, %group0) : (i32, i32, i32, index) -> (i32, i32) + %s1, %m1 = func.call @ggml_q4k_scale_min_from_header(%h0, %h1, %h2, %group1) : (i32, i32, i32, index) -> (i32, i32) + %pair = index.div %group0, %c2 : index + %pair_field = index.mul %pair, %c2 : index + %packet_field = index.div %packet, %c4 : index + %field0 = index.add %pair_field, %packet_field : index + %field = index.add %field0, %c1 : index + %field_add = index.scale %field, %field_bytes : index, offset -> offset + %field_base = index.add %row_base, %field_add : offset + %word_bytes = index.constant 4 : offset + %packet_word = index.rem %packet, %c4 : index + %packet_add = index.scale %packet_word, %word_bytes : index, offset -> offset + %address = index.add %field_base, %packet_add : offset + %codes = func.call @ggml_paired_weight_words(%half_packet, %paired, %is_up, %weight, %peer, %address) : (i1, i1, i1, buffer, buffer, offset) -> (vector<4xi32>) + %mask = vector.constant 252645135 : vector<4xi32> + %shift = vector.constant 4 : vector<4xi32> + %high = vector.shrui %codes, %shift : vector<4xi32> + %nibbles0 = vector.andi %codes, %mask : vector<4xi32> + %bytes0 = vector.bitcast %nibbles0 : vector<4xi32> to vector<16xi8> + %sf0 = scalar.uitofp %s0 : i32 to f32 + %mf0 = scalar.uitofp %m0 : i32 to f32 + %ds0 = scalar.mulf %d, %sf0 : f32 + %ms0 = scalar.mulf %m, %mf0 : f32 + %negative0 = scalar.negf %ms0 : f32 + %dv0 = vector.splat %ds0 : vector<16xf32> + %mv0 = vector.splat %negative0 : vector<16xf32> + %qf0 = vector.uitofp %bytes0 : vector<16xi8> to vector<16xf32> + %values0 = vector.fmaf %qf0, %dv0, %mv0 : vector<16xf32> + %result0 = vector.fptrunc %values0 : vector<16xf32> to vector<16xf16> + %nibbles1 = vector.andi %high, %mask : vector<4xi32> + %bytes1 = vector.bitcast %nibbles1 : vector<4xi32> to vector<16xi8> + %sf1 = scalar.uitofp %s1 : i32 to f32 + %mf1 = scalar.uitofp %m1 : i32 to f32 + %ds1 = scalar.mulf %d, %sf1 : f32 + %ms1 = scalar.mulf %m, %mf1 : f32 + %negative1 = scalar.negf %ms1 : f32 + %dv1 = vector.splat %ds1 : vector<16xf32> + %mv1 = vector.splat %negative1 : vector<16xf32> + %qf1 = vector.uitofp %bytes1 : vector<16xi8> to vector<16xf32> + %values1 = vector.fmaf %qf1, %dv1, %mv1 : vector<16xf32> + %result1 = vector.fptrunc %values1 : vector<16xf32> to vector<16xf16> + 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 + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) { + %zero = vector.constant 0 : vector<2xi32> + %seed = scalar.constant 0 : i32 + %w = vector.bitcast %weight : vector<2xi32> to vector<8xi8> + %a = vector.bitcast %activation : vector<2xi32> to vector<8xi8> + %d = vector.dot4i %w, %a, %zero : vector<8xi8>, vector<8xi8>, vector<2xi32> + %sum = vector.reduce %d, %seed : vector<2xi32>, i32 + %result = scalar.sitofp %sum : i32 to f32 + func.return %result : f32 +} + +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 +func.decl @ggml_q4k_f16_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) + +func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) + +func.decl @ggml_q8_0_f16_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf16>) + +func.decl @ggml_q8_1_f16_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf16>) + +func.decl @ggml_f16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) + +func.decl @ggml_bf16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) + +func.decl @ggml_f32_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_q5k_high_bit_f32(%qh: i8, %qh_bit: i32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c16_i32 = scalar.constant 16 : i32 + %qh_i32 = scalar.extui %qh : i8 to i32 + %masked = scalar.andi %qh_i32, %qh_bit : i32 + %is_set = scalar.cmpi ne, %masked, %c0_i32 : i32 + %high = scf.select %is_set, %c16_i32, %c0_i32 : i32 + %high_f32 = scalar.uitofp %high : i32 to f32 + func.return %high_f32 : f32 +} + +func.def inline @ggml_q5_legacy_high_bit_f32(%qh: i32, %bit: i32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c16_i32 = scalar.constant 16 : i32 + %shifted = scalar.shrui %qh, %bit : i32 + %masked = scalar.andi %shifted, %c1_i32 : i32 + %is_set = scalar.cmpi ne, %masked, %c0_i32 : i32 + %high = scf.select %is_set, %c16_i32, %c0_i32 : i32 + %high_f32 = scalar.uitofp %high : i32 to f32 + func.return %high_f32 : f32 +} + +func.def inline @ggml_q1_0_load_decode_state(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<1xf16>) { + %c0 = index.constant 0 : index + %c128 = index.constant 128 : index + %block_bytes = index.constant 18 : offset + %block = index.div %k, %c128 : index + %block_byte_add = index.scale %block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %d_vector = vector.load %d_view[%c0] : view<1xf16> -> vector<1xf16> + func.return %d_vector : vector<1xf16> +} + +func.def inline @ggml_q1_0_preloaded_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index, %d_vector: vector<1xf16>) -> (vector<4xf32>) { + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c3_i32 = scalar.constant 3 : i32 + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %block_bytes = index.constant 18 : offset + %code_offset = index.constant 2 : offset + %block = index.div %k, %c128 : index + %k_in_block = index.rem %k, %c128 : index + %byte_index0 = index.div %k_in_block, %c8 : index + %byte_index = index.assume %byte_index0 [range(%byte_index0, 0, 15)] : index + %bit_base_index = index.rem %k_in_block, %c8 : index + %bit_base = index.cast %bit_base_index : index to i32 + %block_byte_add = index.scale %block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<16xi8> + %d_f16 = vector.extract %d_vector[0] : vector<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %neg_d = scalar.subf %c0_f32, %d : f32 + %q_i8 = view.load %code_view[%byte_index] : view<16xi8> -> i8 + %q = scalar.extui %q_i8 : i8 to i32 + %bit1 = scalar.addi %bit_base, %c1_i32 : i32 + %bit2 = scalar.addi %bit_base, %c2_i32 : i32 + %bit3 = scalar.addi %bit_base, %c3_i32 : i32 + %q0_shift = scalar.shrui %q, %bit_base : i32 + %q1_shift = scalar.shrui %q, %bit1 : i32 + %q2_shift = scalar.shrui %q, %bit2 : i32 + %q3_shift = scalar.shrui %q, %bit3 : i32 + %q0_mask = scalar.andi %q0_shift, %c1_i32 : i32 + %q1_mask = scalar.andi %q1_shift, %c1_i32 : i32 + %q2_mask = scalar.andi %q2_shift, %c1_i32 : i32 + %q3_mask = scalar.andi %q3_shift, %c1_i32 : i32 + %q0_set = scalar.cmpi ne, %q0_mask, %c0_i32 : i32 + %q1_set = scalar.cmpi ne, %q1_mask, %c0_i32 : i32 + %q2_set = scalar.cmpi ne, %q2_mask, %c0_i32 : i32 + %q3_set = scalar.cmpi ne, %q3_mask, %c0_i32 : i32 + %v0 = scf.select %q0_set, %d, %neg_d : f32 + %v1 = scf.select %q1_set, %d, %neg_d : f32 + %v2 = scf.select %q2_set, %d, %neg_d : f32 + %v3 = scf.select %q3_set, %d, %neg_d : f32 + %result = vector.from_elements %v0, %v1, %v2, %v3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q1_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) { + %d = func.call @ggml_q1_0_load_decode_state(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<1xf16>) + %values = func.call @ggml_q1_0_preloaded_f32_vector4(%weight, %row_byte_base, %k, %d) : (buffer, offset, index, vector<1xf16>) -> (vector<4xf32>) + func.return %values : vector<4xf32> +} + +func.def inline @ggml_q1_0_f16_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q1_0_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +func.def inline @ggml_q4_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c0_i32 = scalar.constant 0 : i32 + %c4_i32 = scalar.constant 4 : i32 + %block_bytes = index.constant 18 : offset + %code_offset = index.constant 2 : offset + %q4_mask = vector.constant 252645135 : vector<1xi32> + %zero_point = vector.constant 8.0 : vector<4xf32> + %block = index.div %k, %c32 : index + %k_in_block = index.rem %k, %c32 : index + %packet = index.div %k_in_block, %c4 : index + %word_index0 = index.rem %packet, %c4 : index + %word_index = index.assume %word_index0 [range(%word_index0, 0, 3)] : index + %is_high = index.cmp uge, %packet, %c4 : index + %shift_i32 = scf.select %is_high, %c4_i32, %c0_i32 : i32 + %shift = vector.splat %shift_i32 : vector<1xi32> + %block_byte_add = index.scale %block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %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<4xi32> + %d_vector = vector.load %d_view[%c0] : view<1xf16> -> vector<1xf16> + %d_f16 = vector.extract %d_vector[0] : vector<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %d_splat = vector.splat %d : vector<4xf32> + %q_word = vector.load %code_view[%word_index] : view<4xi32> -> vector<1xi32> + %shifted_q = vector.shrui %q_word, %shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_unsigned = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %q = vector.subf %q_unsigned, %zero_point : vector<4xf32> + %result = vector.mulf %q, %d_splat : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q4_0_f16_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q4_0_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +func.def inline @ggml_q4_1_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c0_i32 = scalar.constant 0 : i32 + %c4_i32 = scalar.constant 4 : i32 + %block_bytes = index.constant 20 : offset + %code_offset = index.constant 4 : offset + %q4_mask = vector.constant 252645135 : vector<1xi32> + %block = index.div %k, %c32 : index + %k_in_block = index.rem %k, %c32 : index + %packet = index.div %k_in_block, %c4 : index + %word_index0 = index.rem %packet, %c4 : index + %word_index = index.assume %word_index0 [range(%word_index0, 0, 3)] : index + %is_high = index.cmp uge, %packet, %c4 : index + %shift_i32 = scf.select %is_high, %c4_i32, %c0_i32 : i32 + %shift = vector.splat %shift_i32 : vector<1xi32> + %block_byte_add = index.scale %block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<4xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %m_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %m = scalar.extf %m_f16 : f16 to f32 + %d_splat = vector.splat %d : vector<4xf32> + %m_splat = vector.splat %m : vector<4xf32> + %q_word = vector.load %code_view[%word_index] : view<4xi32> -> vector<1xi32> + %shifted_q = vector.shrui %q_word, %shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %scaled = vector.mulf %q, %d_splat : vector<4xf32> + %result = vector.addf %scaled, %m_splat : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q4_1_f16_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q4_1_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +func.def inline @ggml_q5_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c1_i32 = scalar.constant 1 : i32 + %c2 = index.constant 2 : index + %c2_i32 = scalar.constant 2 : i32 + %c4 = index.constant 4 : index + %c16_i32 = scalar.constant 16 : i32 + %c32 = index.constant 32 : index + %c0_i32 = scalar.constant 0 : i32 + %c4_i32 = scalar.constant 4 : i32 + %block_bytes = index.constant 22 : offset + %qh_offset = index.constant 2 : offset + %code_offset = index.constant 6 : offset + %q4_mask = vector.constant 252645135 : vector<1xi32> + %zero_point = vector.constant 16.0 : vector<4xf32> + %block = index.div %k, %c32 : index + %k_in_block = index.rem %k, %c32 : index + %packet = index.div %k_in_block, %c4 : index + %word_index0 = index.rem %packet, %c4 : index + %word_index = index.assume %word_index0 [range(%word_index0, 0, 3)] : index + %is_high = index.cmp uge, %packet, %c4 : index + %shift_i32 = scf.select %is_high, %c4_i32, %c0_i32 : i32 + %high_half_add = scf.select %is_high, %c16_i32, %c0_i32 : i32 + %shift = vector.splat %shift_i32 : vector<1xi32> + %block_byte_add = index.scale %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 + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<1xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<4xi32> + %d_vector = vector.load %d_view[%c0] : view<1xf16> -> vector<1xf16> + %d_f16 = vector.extract %d_vector[0] : vector<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %d_splat = vector.splat %d : vector<4xf32> + %qh_vector = vector.load %qh_view[%c0] : view<1xi32> -> vector<1xi32> + %qh = vector.extract %qh_vector[0] : vector<1xi32> -> i32 + %q_word = vector.load %code_view[%word_index] : view<4xi32> -> vector<1xi32> + %shifted_q = vector.shrui %q_word, %shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_low = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %packet_base = index.mul %word_index, %c4 : index + %packet_base_i32 = index.cast %packet_base : index to i32 + %bit0 = scalar.addi %packet_base_i32, %high_half_add : i32 + %bit1 = scalar.addi %bit0, %c1_i32 : i32 + %bit2 = scalar.addi %bit0, %c2_i32 : i32 + %bit3 = scalar.addi %bit2, %c1_i32 : i32 + %high0 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit0) : (i32, i32) -> (f32) + %high1 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit1) : (i32, i32) -> (f32) + %high2 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit2) : (i32, i32) -> (f32) + %high3 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit3) : (i32, i32) -> (f32) + %high = vector.from_elements %high0, %high1, %high2, %high3 : vector<4xf32> + %q_unsigned = vector.addf %q_low, %high : vector<4xf32> + %q = vector.subf %q_unsigned, %zero_point : vector<4xf32> + %result = vector.mulf %q, %d_splat : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q5_0_f16_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q5_0_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +func.def inline @ggml_q5_1_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c1_i32 = scalar.constant 1 : i32 + %c2 = index.constant 2 : index + %c2_i32 = scalar.constant 2 : i32 + %c4 = index.constant 4 : index + %c16_i32 = scalar.constant 16 : i32 + %c32 = index.constant 32 : index + %c0_i32 = scalar.constant 0 : i32 + %c4_i32 = scalar.constant 4 : i32 + %block_bytes = index.constant 24 : offset + %qh_offset = index.constant 4 : offset + %code_offset = index.constant 8 : offset + %q4_mask = vector.constant 252645135 : vector<1xi32> + %block = index.div %k, %c32 : index + %k_in_block = index.rem %k, %c32 : index + %packet = index.div %k_in_block, %c4 : index + %word_index0 = index.rem %packet, %c4 : index + %word_index = index.assume %word_index0 [range(%word_index0, 0, 3)] : index + %is_high = index.cmp uge, %packet, %c4 : index + %shift_i32 = scf.select %is_high, %c4_i32, %c0_i32 : i32 + %high_half_add = scf.select %is_high, %c16_i32, %c0_i32 : i32 + %shift = vector.splat %shift_i32 : vector<1xi32> + %block_byte_add = index.scale %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 + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<1xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<4xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %m_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %m = scalar.extf %m_f16 : f16 to f32 + %d_splat = vector.splat %d : vector<4xf32> + %m_splat = vector.splat %m : vector<4xf32> + %qh_vector = vector.load %qh_view[%c0] : view<1xi32> -> vector<1xi32> + %qh = vector.extract %qh_vector[0] : vector<1xi32> -> i32 + %q_word = vector.load %code_view[%word_index] : view<4xi32> -> vector<1xi32> + %shifted_q = vector.shrui %q_word, %shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_low = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %packet_base = index.mul %word_index, %c4 : index + %packet_base_i32 = index.cast %packet_base : index to i32 + %bit0 = scalar.addi %packet_base_i32, %high_half_add : i32 + %bit1 = scalar.addi %bit0, %c1_i32 : i32 + %bit2 = scalar.addi %bit0, %c2_i32 : i32 + %bit3 = scalar.addi %bit2, %c1_i32 : i32 + %high0 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit0) : (i32, i32) -> (f32) + %high1 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit1) : (i32, i32) -> (f32) + %high2 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit2) : (i32, i32) -> (f32) + %high3 = func.call @ggml_q5_legacy_high_bit_f32(%qh, %bit3) : (i32, i32) -> (f32) + %high = vector.from_elements %high0, %high1, %high2, %high3 : vector<4xf32> + %q = vector.addf %q_low, %high : vector<4xf32> + %scaled = vector.mulf %q, %d_splat : vector<4xf32> + %result = vector.addf %scaled, %m_splat : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q5_1_f16_vector4(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q5_1_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +// Decodes four adjacent Q4_K values to F32 without the FP16 staging round used +// by WMMA paths. +func.def inline @ggml_q4k_f32_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c0_i32 = scalar.constant 0 : i32 + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %scale_offset = index.constant 4 : offset + %code_offset = index.constant 16 : offset + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<3xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +// Decode the byte-identical row64 native Q4_K carrier used by low-token +// contractions. The header and paired-nibble fields retain GGUF values. +func.def inline @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_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 + %block_bytes = index.constant 144 : offset + %scale_offset = index.constant 4 : offset + %code_offset = index.constant 16 : offset + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %record_bytes = index.constant 9216 : offset + %row_bytes = index.constant 16 : offset + %payload_offset = index.constant 1024 : offset + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %group_block = index.madd %row_group, %blocks, %q4_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %row_add = index.scale %row_lane, %row_bytes : index, offset -> offset + %block_byte_base = index.add %group_base, %row_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %code_byte_base = index.add %group_base, %payload_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<3xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<8x64x4xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %payload_field = index.div %q_word_index, %c4 : index + %payload_word = index.rem %q_word_index, %c4 : index + %q_word = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +// Q5_K uses the Q4_K affine header and a 32-byte fifth-bit plane. +func.def inline @ggml_q5k_f32_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_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 + %block_bytes = index.constant 176 : offset + %scale_offset = index.constant 4 : offset + %qh_offset = index.constant 16 : offset + %code_offset = index.constant 48 : offset + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q5_group [range(%q5_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q5_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %qh_byte_base = index.add %block_byte_base, %qh_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<3xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<32xi8> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %packet_base = index.mul %bounded_packet, %c4 : index + %qh_bytes = vector.load %qh_view[%packet_base] : view<32xi8> -> vector<4xi8> + %qh_shift_i32 = index.cast %bounded_group : index to i32 + %qh_bit = scalar.shli %c1_i32, %qh_shift_i32 : i32 + %qh0 = vector.extract %qh_bytes[0] : vector<4xi8> -> i8 + %qh1 = vector.extract %qh_bytes[1] : vector<4xi8> -> i8 + %qh2 = vector.extract %qh_bytes[2] : vector<4xi8> -> i8 + %qh3 = vector.extract %qh_bytes[3] : vector<4xi8> -> i8 + %high0 = func.call @ggml_q5k_high_bit_f32(%qh0, %qh_bit) : (i8, i32) -> (f32) + %high1 = func.call @ggml_q5k_high_bit_f32(%qh1, %qh_bit) : (i8, i32) -> (f32) + %high2 = func.call @ggml_q5k_high_bit_f32(%qh2, %qh_bit) : (i8, i32) -> (f32) + %high3 = func.call @ggml_q5k_high_bit_f32(%qh3, %qh_bit) : (i8, i32) -> (f32) + %q0_low = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1_low = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2_low = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3_low = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %q0 = scalar.addf %q0_low, %high0 : f32 + %q1 = scalar.addf %q1_low, %high1 : f32 + %q2 = scalar.addf %q2_low, %high2 : f32 + %q3 = scalar.addf %q3_low, %high3 : f32 + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q5k_f16_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q5k_f32_vector4(%weight, %row_byte_base, %q5_block, %q5_group, %packet) : (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_q3k_value_f32(%q_byte: i8, %hmask_byte: i8, %q_shift: i32, %hmask_bit: i32, %combined_scale: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c3_i32 = scalar.constant 3 : i32 + %c4_i32 = scalar.constant 4 : i32 + %q_i32 = scalar.extui %q_byte : i8 to i32 + %q_shifted = scalar.shrui %q_i32, %q_shift : i32 + %q_low = scalar.andi %q_shifted, %c3_i32 : i32 + %hmask_i32 = scalar.extui %hmask_byte : i8 to i32 + %hmask_masked = scalar.andi %hmask_i32, %hmask_bit : i32 + %has_high_bit = scalar.cmpi ne, %hmask_masked, %c0_i32 : i32 + %subtrahend = scf.select %has_high_bit, %c0_i32, %c4_i32 : i32 + %q = scalar.subi %q_low, %subtrahend : i32 + %q_f32 = scalar.sitofp %q : i32 to f32 + %value = scalar.mulf %q_f32, %combined_scale : f32 + func.return %value : f32 +} + +func.def inline @ggml_q3k_f32_vector4(%weight: buffer, %row_byte_base: offset, %q3_block: index, %q3_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 + %c32 = index.constant 32 : index + %block_bytes = index.constant 110 : offset + %code_offset = index.constant 32 : offset + %scale_offset = index.constant 96 : offset + %d_offset = index.constant 108 : 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 + %c15_i32 = scalar.constant 15 : i32 + %c32_i32 = scalar.constant 32 : i32 + %bounded_group = index.assume %q3_group [range(%q3_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q3_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %d_byte_base = index.add %block_byte_base, %d_offset : offset + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %hmask_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi8> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<64xi8> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<12xi8> + %scale_subgroup = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index0 = index.add %scale_group_base, %scale_subgroup : index + %scale_index = index.assume %scale_index0 [range(%scale_index0, 0, 15)] : index + %uses_high_nibble = index.cmp uge, %scale_index, %c8 : index + %scale_low_index0 = index.rem %scale_index, %c8 : index + %scale_low_index = index.assume %scale_low_index0 [range(%scale_low_index0, 0, 7)] : index + %scale_high_index_rem = index.rem %scale_index, %c4 : index + %scale_high_index0 = index.add %scale_high_index_rem, %c8 : index + %scale_high_index = index.assume %scale_high_index0 [range(%scale_high_index0, 8, 11)] : index + %scale_low_byte = view.load %scale_view[%scale_low_index] : view<12xi8> -> i8 + %scale_high_byte = view.load %scale_view[%scale_high_index] : view<12xi8> -> i8 + %scale_low_i32 = scalar.extui %scale_low_byte : i8 to i32 + %scale_high_i32 = scalar.extui %scale_high_byte : i8 to i32 + %scale_low_shift = scf.select %uses_high_nibble, %c4_i32, %c0_i32 : i32 + %scale_low_shifted = scalar.shrui %scale_low_i32, %scale_low_shift : i32 + %scale_low = scalar.andi %scale_low_shifted, %c15_i32 : i32 + %scale_group_quartile = index.div %scale_index, %c4 : index + %scale_high_shift_index = index.mul %scale_group_quartile, %c2 : index + %scale_high_shift = index.cast %scale_high_shift_index : index to i32 + %scale_high_shifted = scalar.shrui %scale_high_i32, %scale_high_shift : i32 + %scale_high_low = scalar.andi %scale_high_shifted, %c3_i32 : i32 + %scale_high = scalar.shli %scale_high_low, %c4_i32 : i32 + %scale_code = scalar.ori %scale_low, %scale_high : i32 + %centered_scale = scalar.subi %scale_code, %c32_i32 : i32 + %centered_scale_f32 = scalar.sitofp %centered_scale : i32 to f32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %d, %centered_scale_f32 : f32 + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %packet_base = index.mul %bounded_packet, %c4 : index + %code_half_base = index.mul %half, %c32 : index + %code_index0 = index.add %code_half_base, %packet_base : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 63)] : index + %hmask_index = index.assume %packet_base [range(%packet_base, 0, 28), mul(%packet_base, 4)] : index + %q_shift_index = index.mul %group_in_half, %c2 : index + %q_shift = index.cast %q_shift_index : index to i32 + %hmask_bit_shift = index.cast %bounded_group : index to i32 + %hmask_bit = scalar.shli %c1_i32, %hmask_bit_shift : i32 + %q_bytes = vector.load %code_view[%code_index] : view<64xi8> -> vector<4xi8> + %hmask_bytes = vector.load %hmask_view[%hmask_index] : view<32xi8> -> vector<4xi8> + %q0 = vector.extract %q_bytes[0] : vector<4xi8> -> i8 + %q1 = vector.extract %q_bytes[1] : vector<4xi8> -> i8 + %q2 = vector.extract %q_bytes[2] : vector<4xi8> -> i8 + %q3 = vector.extract %q_bytes[3] : vector<4xi8> -> i8 + %hm0 = vector.extract %hmask_bytes[0] : vector<4xi8> -> i8 + %hm1 = vector.extract %hmask_bytes[1] : vector<4xi8> -> i8 + %hm2 = vector.extract %hmask_bytes[2] : vector<4xi8> -> i8 + %hm3 = vector.extract %hmask_bytes[3] : vector<4xi8> -> i8 + %value0 = func.call @ggml_q3k_value_f32(%q0, %hm0, %q_shift, %hmask_bit, %combined_scale) : (i8, i8, i32, i32, f32) -> (f32) + %value1 = func.call @ggml_q3k_value_f32(%q1, %hm1, %q_shift, %hmask_bit, %combined_scale) : (i8, i8, i32, i32, f32) -> (f32) + %value2 = func.call @ggml_q3k_value_f32(%q2, %hm2, %q_shift, %hmask_bit, %combined_scale) : (i8, i8, i32, i32, f32) -> (f32) + %value3 = func.call @ggml_q3k_value_f32(%q3, %hm3, %q_shift, %hmask_bit, %combined_scale) : (i8, i8, i32, i32, f32) -> (f32) + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_q3k_f16_vector4(%weight: buffer, %row_byte_base: offset, %q3_block: index, %q3_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_q3k_f32_vector4(%weight, %row_byte_base, %q3_block, %q3_group, %packet) : (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_iq2s_grid_lookup_i32(%grid_index: i32, %word_index: i32) -> (i32) { + %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 + %c9_i32 = scalar.constant 9 : i32 + %c10_i32 = scalar.constant 10 : i32 + %c11_i32 = scalar.constant 11 : i32 + %c12_i32 = scalar.constant 12 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c14_i32 = scalar.constant 14 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c17_i32 = scalar.constant 17 : i32 + %c18_i32 = scalar.constant 18 : i32 + %c19_i32 = scalar.constant 19 : i32 + %c20_i32 = scalar.constant 20 : i32 + %c21_i32 = scalar.constant 21 : i32 + %c22_i32 = scalar.constant 22 : i32 + %c23_i32 = scalar.constant 23 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c25_i32 = scalar.constant 25 : i32 + %c26_i32 = scalar.constant 26 : i32 + %c27_i32 = scalar.constant 27 : i32 + %c28_i32 = scalar.constant 28 : i32 + %c29_i32 = scalar.constant 29 : i32 + %c30_i32 = scalar.constant 30 : i32 + %c31_i32 = scalar.constant 31 : i32 + %chunk_i32 = scalar.shrui %grid_index, %c5_i32 : i32 + %lane_i32 = scalar.andi %grid_index, %c31_i32 : i32 + %codes = vector.from_elements %lane_i32 : vector<1xi32> + %is_chunk1 = scalar.cmpi eq, %chunk_i32, %c1_i32 : i32 + %is_chunk2 = scalar.cmpi eq, %chunk_i32, %c2_i32 : i32 + %is_chunk3 = scalar.cmpi eq, %chunk_i32, %c3_i32 : i32 + %is_chunk4 = scalar.cmpi eq, %chunk_i32, %c4_i32 : i32 + %is_chunk5 = scalar.cmpi eq, %chunk_i32, %c5_i32 : i32 + %is_chunk6 = scalar.cmpi eq, %chunk_i32, %c6_i32 : i32 + %is_chunk7 = scalar.cmpi eq, %chunk_i32, %c7_i32 : i32 + %is_chunk8 = scalar.cmpi eq, %chunk_i32, %c8_i32 : i32 + %is_chunk9 = scalar.cmpi eq, %chunk_i32, %c9_i32 : i32 + %is_chunk10 = scalar.cmpi eq, %chunk_i32, %c10_i32 : i32 + %is_chunk11 = scalar.cmpi eq, %chunk_i32, %c11_i32 : i32 + %is_chunk12 = scalar.cmpi eq, %chunk_i32, %c12_i32 : i32 + %is_chunk13 = scalar.cmpi eq, %chunk_i32, %c13_i32 : i32 + %is_chunk14 = scalar.cmpi eq, %chunk_i32, %c14_i32 : i32 + %is_chunk15 = scalar.cmpi eq, %chunk_i32, %c15_i32 : i32 + %is_chunk16 = scalar.cmpi eq, %chunk_i32, %c16_i32 : i32 + %is_chunk17 = scalar.cmpi eq, %chunk_i32, %c17_i32 : i32 + %is_chunk18 = scalar.cmpi eq, %chunk_i32, %c18_i32 : i32 + %is_chunk19 = scalar.cmpi eq, %chunk_i32, %c19_i32 : i32 + %is_chunk20 = scalar.cmpi eq, %chunk_i32, %c20_i32 : i32 + %is_chunk21 = scalar.cmpi eq, %chunk_i32, %c21_i32 : i32 + %is_chunk22 = scalar.cmpi eq, %chunk_i32, %c22_i32 : i32 + %is_chunk23 = scalar.cmpi eq, %chunk_i32, %c23_i32 : i32 + %is_chunk24 = scalar.cmpi eq, %chunk_i32, %c24_i32 : i32 + %is_chunk25 = scalar.cmpi eq, %chunk_i32, %c25_i32 : i32 + %is_chunk26 = scalar.cmpi eq, %chunk_i32, %c26_i32 : i32 + %is_chunk27 = scalar.cmpi eq, %chunk_i32, %c27_i32 : i32 + %is_chunk28 = scalar.cmpi eq, %chunk_i32, %c28_i32 : i32 + %is_chunk29 = scalar.cmpi eq, %chunk_i32, %c29_i32 : i32 + %is_chunk30 = scalar.cmpi eq, %chunk_i32, %c30_i32 : i32 + %is_chunk31 = scalar.cmpi eq, %chunk_i32, %c31_i32 : i32 + %iq2s_word0_low_0_0 = scalar.constant 2056.0 : f32 + %iq2s_word0_low_0_1 = scalar.constant 2091.0 : f32 + %iq2s_word0_low_0_2 = scalar.constant 6425.0 : f32 + %iq2s_word0_low_0_3 = scalar.constant 11016.0 : f32 + %iq2s_word0_low_0_4 = scalar.constant 11051.0 : f32 + %iq2s_word0_low_0_5 = scalar.constant 2073.0 : f32 + %iq2s_word0_low_0_6 = scalar.constant 6408.0 : f32 + %iq2s_word0_low_0_7 = scalar.constant 6443.0 : f32 + %iq2s_word0_low_0_8 = scalar.constant 11033.0 : f32 + %iq2s_word0_low_0 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_low_1 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_low_2 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_6, %iq2s_word0_low_0_1, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_low_3 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_low_4 = vector.from_elements %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_low_5 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_6, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_low_6 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word0_low_7 = vector.from_elements %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_low_8 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8 : vector<32xf32> + %iq2s_word0_low_9 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_low_10 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_low_11 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_low_12 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_6, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word0_low_13 = vector.from_elements %iq2s_word0_low_0_4, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_low_14 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word0_low_15 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_low_16 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_low_17 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_low_18 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_low_19 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_low_20 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_low_21 = vector.from_elements %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word0_low_22 = vector.from_elements %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word0_low_23 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_low_24 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_low_25 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_3, %iq2s_word0_low_0_1, %iq2s_word0_low_0_0, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_1, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_low_26 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_low_27 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_low_28 = vector.from_elements %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_low_29 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_low_30 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_6, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_1, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_6, %iq2s_word0_low_0_1, %iq2s_word0_low_0_4, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word0_low_31 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_2, %iq2s_word0_low_0_1, %iq2s_word0_low_0_4, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_3, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_5, %iq2s_word0_low_0_4, %iq2s_word0_low_0_3, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4 : vector<32xf32> + %selected_word0_low1 = scf.select %is_chunk1, %iq2s_word0_low_1, %iq2s_word0_low_0 : vector<32xf32> + %selected_word0_low2 = scf.select %is_chunk2, %iq2s_word0_low_2, %selected_word0_low1 : vector<32xf32> + %selected_word0_low3 = scf.select %is_chunk3, %iq2s_word0_low_3, %selected_word0_low2 : vector<32xf32> + %selected_word0_low4 = scf.select %is_chunk4, %iq2s_word0_low_4, %selected_word0_low3 : vector<32xf32> + %selected_word0_low5 = scf.select %is_chunk5, %iq2s_word0_low_5, %selected_word0_low4 : vector<32xf32> + %selected_word0_low6 = scf.select %is_chunk6, %iq2s_word0_low_6, %selected_word0_low5 : vector<32xf32> + %selected_word0_low7 = scf.select %is_chunk7, %iq2s_word0_low_7, %selected_word0_low6 : vector<32xf32> + %selected_word0_low8 = scf.select %is_chunk8, %iq2s_word0_low_8, %selected_word0_low7 : vector<32xf32> + %selected_word0_low9 = scf.select %is_chunk9, %iq2s_word0_low_9, %selected_word0_low8 : vector<32xf32> + %selected_word0_low10 = scf.select %is_chunk10, %iq2s_word0_low_10, %selected_word0_low9 : vector<32xf32> + %selected_word0_low11 = scf.select %is_chunk11, %iq2s_word0_low_11, %selected_word0_low10 : vector<32xf32> + %selected_word0_low12 = scf.select %is_chunk12, %iq2s_word0_low_12, %selected_word0_low11 : vector<32xf32> + %selected_word0_low13 = scf.select %is_chunk13, %iq2s_word0_low_13, %selected_word0_low12 : vector<32xf32> + %selected_word0_low14 = scf.select %is_chunk14, %iq2s_word0_low_14, %selected_word0_low13 : vector<32xf32> + %selected_word0_low15 = scf.select %is_chunk15, %iq2s_word0_low_15, %selected_word0_low14 : vector<32xf32> + %selected_word0_low16 = scf.select %is_chunk16, %iq2s_word0_low_16, %selected_word0_low15 : vector<32xf32> + %selected_word0_low17 = scf.select %is_chunk17, %iq2s_word0_low_17, %selected_word0_low16 : vector<32xf32> + %selected_word0_low18 = scf.select %is_chunk18, %iq2s_word0_low_18, %selected_word0_low17 : vector<32xf32> + %selected_word0_low19 = scf.select %is_chunk19, %iq2s_word0_low_19, %selected_word0_low18 : vector<32xf32> + %selected_word0_low20 = scf.select %is_chunk20, %iq2s_word0_low_20, %selected_word0_low19 : vector<32xf32> + %selected_word0_low21 = scf.select %is_chunk21, %iq2s_word0_low_21, %selected_word0_low20 : vector<32xf32> + %selected_word0_low22 = scf.select %is_chunk22, %iq2s_word0_low_22, %selected_word0_low21 : vector<32xf32> + %selected_word0_low23 = scf.select %is_chunk23, %iq2s_word0_low_23, %selected_word0_low22 : vector<32xf32> + %selected_word0_low24 = scf.select %is_chunk24, %iq2s_word0_low_24, %selected_word0_low23 : vector<32xf32> + %selected_word0_low25 = scf.select %is_chunk25, %iq2s_word0_low_25, %selected_word0_low24 : vector<32xf32> + %selected_word0_low26 = scf.select %is_chunk26, %iq2s_word0_low_26, %selected_word0_low25 : vector<32xf32> + %selected_word0_low27 = scf.select %is_chunk27, %iq2s_word0_low_27, %selected_word0_low26 : vector<32xf32> + %selected_word0_low28 = scf.select %is_chunk28, %iq2s_word0_low_28, %selected_word0_low27 : vector<32xf32> + %selected_word0_low29 = scf.select %is_chunk29, %iq2s_word0_low_29, %selected_word0_low28 : vector<32xf32> + %selected_word0_low30 = scf.select %is_chunk30, %iq2s_word0_low_30, %selected_word0_low29 : vector<32xf32> + %selected_word0_low31 = scf.select %is_chunk31, %iq2s_word0_low_31, %selected_word0_low30 : vector<32xf32> + %iq2s_word0_high_0 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word0_high_1 = vector.from_elements %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_high_2 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_3 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_4 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_high_5 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_6 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_high_7 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_high_8 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_high_9 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_10 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_high_11 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_12 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word0_high_13 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word0_high_14 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word0_high_15 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_high_16 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word0_high_17 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_high_18 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word0_high_19 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_high_20 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_high_21 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word0_high_22 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word0_high_23 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_24 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_high_25 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_7, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_high_26 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_27 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word0_high_28 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_1, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word0_high_29 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word0_high_30 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_8, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word0_high_31 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_7, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_6, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_2, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_4, %iq2s_word0_low_0_2, %iq2s_word0_low_0_3, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_3, %iq2s_word0_low_0_4 : vector<32xf32> + %selected_word0_high1 = scf.select %is_chunk1, %iq2s_word0_high_1, %iq2s_word0_high_0 : vector<32xf32> + %selected_word0_high2 = scf.select %is_chunk2, %iq2s_word0_high_2, %selected_word0_high1 : vector<32xf32> + %selected_word0_high3 = scf.select %is_chunk3, %iq2s_word0_high_3, %selected_word0_high2 : vector<32xf32> + %selected_word0_high4 = scf.select %is_chunk4, %iq2s_word0_high_4, %selected_word0_high3 : vector<32xf32> + %selected_word0_high5 = scf.select %is_chunk5, %iq2s_word0_high_5, %selected_word0_high4 : vector<32xf32> + %selected_word0_high6 = scf.select %is_chunk6, %iq2s_word0_high_6, %selected_word0_high5 : vector<32xf32> + %selected_word0_high7 = scf.select %is_chunk7, %iq2s_word0_high_7, %selected_word0_high6 : vector<32xf32> + %selected_word0_high8 = scf.select %is_chunk8, %iq2s_word0_high_8, %selected_word0_high7 : vector<32xf32> + %selected_word0_high9 = scf.select %is_chunk9, %iq2s_word0_high_9, %selected_word0_high8 : vector<32xf32> + %selected_word0_high10 = scf.select %is_chunk10, %iq2s_word0_high_10, %selected_word0_high9 : vector<32xf32> + %selected_word0_high11 = scf.select %is_chunk11, %iq2s_word0_high_11, %selected_word0_high10 : vector<32xf32> + %selected_word0_high12 = scf.select %is_chunk12, %iq2s_word0_high_12, %selected_word0_high11 : vector<32xf32> + %selected_word0_high13 = scf.select %is_chunk13, %iq2s_word0_high_13, %selected_word0_high12 : vector<32xf32> + %selected_word0_high14 = scf.select %is_chunk14, %iq2s_word0_high_14, %selected_word0_high13 : vector<32xf32> + %selected_word0_high15 = scf.select %is_chunk15, %iq2s_word0_high_15, %selected_word0_high14 : vector<32xf32> + %selected_word0_high16 = scf.select %is_chunk16, %iq2s_word0_high_16, %selected_word0_high15 : vector<32xf32> + %selected_word0_high17 = scf.select %is_chunk17, %iq2s_word0_high_17, %selected_word0_high16 : vector<32xf32> + %selected_word0_high18 = scf.select %is_chunk18, %iq2s_word0_high_18, %selected_word0_high17 : vector<32xf32> + %selected_word0_high19 = scf.select %is_chunk19, %iq2s_word0_high_19, %selected_word0_high18 : vector<32xf32> + %selected_word0_high20 = scf.select %is_chunk20, %iq2s_word0_high_20, %selected_word0_high19 : vector<32xf32> + %selected_word0_high21 = scf.select %is_chunk21, %iq2s_word0_high_21, %selected_word0_high20 : vector<32xf32> + %selected_word0_high22 = scf.select %is_chunk22, %iq2s_word0_high_22, %selected_word0_high21 : vector<32xf32> + %selected_word0_high23 = scf.select %is_chunk23, %iq2s_word0_high_23, %selected_word0_high22 : vector<32xf32> + %selected_word0_high24 = scf.select %is_chunk24, %iq2s_word0_high_24, %selected_word0_high23 : vector<32xf32> + %selected_word0_high25 = scf.select %is_chunk25, %iq2s_word0_high_25, %selected_word0_high24 : vector<32xf32> + %selected_word0_high26 = scf.select %is_chunk26, %iq2s_word0_high_26, %selected_word0_high25 : vector<32xf32> + %selected_word0_high27 = scf.select %is_chunk27, %iq2s_word0_high_27, %selected_word0_high26 : vector<32xf32> + %selected_word0_high28 = scf.select %is_chunk28, %iq2s_word0_high_28, %selected_word0_high27 : vector<32xf32> + %selected_word0_high29 = scf.select %is_chunk29, %iq2s_word0_high_29, %selected_word0_high28 : vector<32xf32> + %selected_word0_high30 = scf.select %is_chunk30, %iq2s_word0_high_30, %selected_word0_high29 : vector<32xf32> + %selected_word0_high31 = scf.select %is_chunk31, %iq2s_word0_high_31, %selected_word0_high30 : vector<32xf32> + %iq2s_word1_low_0 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_1 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_low_2 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_low_3 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_low_4 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word1_low_5 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8 : vector<32xf32> + %iq2s_word1_low_6 = vector.from_elements %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_7 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_low_8 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_low_9 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_low_10 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word1_low_11 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_12 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_low_13 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_low_14 = vector.from_elements %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_15 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_low_16 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_low_17 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_low_18 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word1_low_19 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word1_low_20 = vector.from_elements %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_low_21 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_low_22 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word1_low_23 = vector.from_elements %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_24 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_low_25 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_26 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_low_27 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_low_28 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_29 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_low_30 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_low_31 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4 : vector<32xf32> + %selected_word1_low1 = scf.select %is_chunk1, %iq2s_word1_low_1, %iq2s_word1_low_0 : vector<32xf32> + %selected_word1_low2 = scf.select %is_chunk2, %iq2s_word1_low_2, %selected_word1_low1 : vector<32xf32> + %selected_word1_low3 = scf.select %is_chunk3, %iq2s_word1_low_3, %selected_word1_low2 : vector<32xf32> + %selected_word1_low4 = scf.select %is_chunk4, %iq2s_word1_low_4, %selected_word1_low3 : vector<32xf32> + %selected_word1_low5 = scf.select %is_chunk5, %iq2s_word1_low_5, %selected_word1_low4 : vector<32xf32> + %selected_word1_low6 = scf.select %is_chunk6, %iq2s_word1_low_6, %selected_word1_low5 : vector<32xf32> + %selected_word1_low7 = scf.select %is_chunk7, %iq2s_word1_low_7, %selected_word1_low6 : vector<32xf32> + %selected_word1_low8 = scf.select %is_chunk8, %iq2s_word1_low_8, %selected_word1_low7 : vector<32xf32> + %selected_word1_low9 = scf.select %is_chunk9, %iq2s_word1_low_9, %selected_word1_low8 : vector<32xf32> + %selected_word1_low10 = scf.select %is_chunk10, %iq2s_word1_low_10, %selected_word1_low9 : vector<32xf32> + %selected_word1_low11 = scf.select %is_chunk11, %iq2s_word1_low_11, %selected_word1_low10 : vector<32xf32> + %selected_word1_low12 = scf.select %is_chunk12, %iq2s_word1_low_12, %selected_word1_low11 : vector<32xf32> + %selected_word1_low13 = scf.select %is_chunk13, %iq2s_word1_low_13, %selected_word1_low12 : vector<32xf32> + %selected_word1_low14 = scf.select %is_chunk14, %iq2s_word1_low_14, %selected_word1_low13 : vector<32xf32> + %selected_word1_low15 = scf.select %is_chunk15, %iq2s_word1_low_15, %selected_word1_low14 : vector<32xf32> + %selected_word1_low16 = scf.select %is_chunk16, %iq2s_word1_low_16, %selected_word1_low15 : vector<32xf32> + %selected_word1_low17 = scf.select %is_chunk17, %iq2s_word1_low_17, %selected_word1_low16 : vector<32xf32> + %selected_word1_low18 = scf.select %is_chunk18, %iq2s_word1_low_18, %selected_word1_low17 : vector<32xf32> + %selected_word1_low19 = scf.select %is_chunk19, %iq2s_word1_low_19, %selected_word1_low18 : vector<32xf32> + %selected_word1_low20 = scf.select %is_chunk20, %iq2s_word1_low_20, %selected_word1_low19 : vector<32xf32> + %selected_word1_low21 = scf.select %is_chunk21, %iq2s_word1_low_21, %selected_word1_low20 : vector<32xf32> + %selected_word1_low22 = scf.select %is_chunk22, %iq2s_word1_low_22, %selected_word1_low21 : vector<32xf32> + %selected_word1_low23 = scf.select %is_chunk23, %iq2s_word1_low_23, %selected_word1_low22 : vector<32xf32> + %selected_word1_low24 = scf.select %is_chunk24, %iq2s_word1_low_24, %selected_word1_low23 : vector<32xf32> + %selected_word1_low25 = scf.select %is_chunk25, %iq2s_word1_low_25, %selected_word1_low24 : vector<32xf32> + %selected_word1_low26 = scf.select %is_chunk26, %iq2s_word1_low_26, %selected_word1_low25 : vector<32xf32> + %selected_word1_low27 = scf.select %is_chunk27, %iq2s_word1_low_27, %selected_word1_low26 : vector<32xf32> + %selected_word1_low28 = scf.select %is_chunk28, %iq2s_word1_low_28, %selected_word1_low27 : vector<32xf32> + %selected_word1_low29 = scf.select %is_chunk29, %iq2s_word1_low_29, %selected_word1_low28 : vector<32xf32> + %selected_word1_low30 = scf.select %is_chunk30, %iq2s_word1_low_30, %selected_word1_low29 : vector<32xf32> + %selected_word1_low31 = scf.select %is_chunk31, %iq2s_word1_low_31, %selected_word1_low30 : vector<32xf32> + %iq2s_word1_high_0 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_high_1 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_high_2 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_high_3 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_high_4 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_high_5 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0 : vector<32xf32> + %iq2s_word1_high_6 = vector.from_elements %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_0, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_high_7 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_high_8 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_high_9 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_high_10 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5 : vector<32xf32> + %iq2s_word1_high_11 = vector.from_elements %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_5, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_high_12 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_high_13 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1 : vector<32xf32> + %iq2s_word1_high_14 = vector.from_elements %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_1, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_high_15 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_high_16 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_high_17 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_high_18 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_high_19 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_6 : vector<32xf32> + %iq2s_word1_high_20 = vector.from_elements %iq2s_word0_low_0_6, %iq2s_word0_low_0_6, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_high_21 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_high_22 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2 : vector<32xf32> + %iq2s_word1_high_23 = vector.from_elements %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_2, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word1_high_24 = vector.from_elements %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7 : vector<32xf32> + %iq2s_word1_high_25 = vector.from_elements %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_7, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word1_high_26 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word1_high_27 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3 : vector<32xf32> + %iq2s_word1_high_28 = vector.from_elements %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_3, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8 : vector<32xf32> + %iq2s_word1_high_29 = vector.from_elements %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8 : vector<32xf32> + %iq2s_word1_high_30 = vector.from_elements %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_8, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4 : vector<32xf32> + %iq2s_word1_high_31 = vector.from_elements %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4, %iq2s_word0_low_0_4 : vector<32xf32> + %selected_word1_high1 = scf.select %is_chunk1, %iq2s_word1_high_1, %iq2s_word1_high_0 : vector<32xf32> + %selected_word1_high2 = scf.select %is_chunk2, %iq2s_word1_high_2, %selected_word1_high1 : vector<32xf32> + %selected_word1_high3 = scf.select %is_chunk3, %iq2s_word1_high_3, %selected_word1_high2 : vector<32xf32> + %selected_word1_high4 = scf.select %is_chunk4, %iq2s_word1_high_4, %selected_word1_high3 : vector<32xf32> + %selected_word1_high5 = scf.select %is_chunk5, %iq2s_word1_high_5, %selected_word1_high4 : vector<32xf32> + %selected_word1_high6 = scf.select %is_chunk6, %iq2s_word1_high_6, %selected_word1_high5 : vector<32xf32> + %selected_word1_high7 = scf.select %is_chunk7, %iq2s_word1_high_7, %selected_word1_high6 : vector<32xf32> + %selected_word1_high8 = scf.select %is_chunk8, %iq2s_word1_high_8, %selected_word1_high7 : vector<32xf32> + %selected_word1_high9 = scf.select %is_chunk9, %iq2s_word1_high_9, %selected_word1_high8 : vector<32xf32> + %selected_word1_high10 = scf.select %is_chunk10, %iq2s_word1_high_10, %selected_word1_high9 : vector<32xf32> + %selected_word1_high11 = scf.select %is_chunk11, %iq2s_word1_high_11, %selected_word1_high10 : vector<32xf32> + %selected_word1_high12 = scf.select %is_chunk12, %iq2s_word1_high_12, %selected_word1_high11 : vector<32xf32> + %selected_word1_high13 = scf.select %is_chunk13, %iq2s_word1_high_13, %selected_word1_high12 : vector<32xf32> + %selected_word1_high14 = scf.select %is_chunk14, %iq2s_word1_high_14, %selected_word1_high13 : vector<32xf32> + %selected_word1_high15 = scf.select %is_chunk15, %iq2s_word1_high_15, %selected_word1_high14 : vector<32xf32> + %selected_word1_high16 = scf.select %is_chunk16, %iq2s_word1_high_16, %selected_word1_high15 : vector<32xf32> + %selected_word1_high17 = scf.select %is_chunk17, %iq2s_word1_high_17, %selected_word1_high16 : vector<32xf32> + %selected_word1_high18 = scf.select %is_chunk18, %iq2s_word1_high_18, %selected_word1_high17 : vector<32xf32> + %selected_word1_high19 = scf.select %is_chunk19, %iq2s_word1_high_19, %selected_word1_high18 : vector<32xf32> + %selected_word1_high20 = scf.select %is_chunk20, %iq2s_word1_high_20, %selected_word1_high19 : vector<32xf32> + %selected_word1_high21 = scf.select %is_chunk21, %iq2s_word1_high_21, %selected_word1_high20 : vector<32xf32> + %selected_word1_high22 = scf.select %is_chunk22, %iq2s_word1_high_22, %selected_word1_high21 : vector<32xf32> + %selected_word1_high23 = scf.select %is_chunk23, %iq2s_word1_high_23, %selected_word1_high22 : vector<32xf32> + %selected_word1_high24 = scf.select %is_chunk24, %iq2s_word1_high_24, %selected_word1_high23 : vector<32xf32> + %selected_word1_high25 = scf.select %is_chunk25, %iq2s_word1_high_25, %selected_word1_high24 : vector<32xf32> + %selected_word1_high26 = scf.select %is_chunk26, %iq2s_word1_high_26, %selected_word1_high25 : vector<32xf32> + %selected_word1_high27 = scf.select %is_chunk27, %iq2s_word1_high_27, %selected_word1_high26 : vector<32xf32> + %selected_word1_high28 = scf.select %is_chunk28, %iq2s_word1_high_28, %selected_word1_high27 : vector<32xf32> + %selected_word1_high29 = scf.select %is_chunk29, %iq2s_word1_high_29, %selected_word1_high28 : vector<32xf32> + %selected_word1_high30 = scf.select %is_chunk30, %iq2s_word1_high_30, %selected_word1_high29 : vector<32xf32> + %selected_word1_high31 = scf.select %is_chunk31, %iq2s_word1_high_31, %selected_word1_high30 : vector<32xf32> + %use_word1 = scalar.cmpi eq, %word_index, %c1_i32 : i32 + %selected_low = scf.select %use_word1, %selected_word1_low31, %selected_word0_low31 : vector<32xf32> + %selected_high = scf.select %use_word1, %selected_word1_high31, %selected_word0_high31 : vector<32xf32> + %low_values = vector.table.lookup %selected_low[%codes] : vector<32xf32>, vector<1xi32> -> vector<1xf32> + %high_values = vector.table.lookup %selected_high[%codes] : vector<32xf32>, vector<1xi32> -> vector<1xf32> + %low_f32 = vector.extract %low_values[0] : vector<1xf32> -> f32 + %high_f32 = vector.extract %high_values[0] : vector<1xf32> -> f32 + %low = scalar.fptoui %low_f32 : f32 to i32 + %high = scalar.fptoui %high_f32 : f32 to i32 + %high_shifted = scalar.shli %high, %c16_i32 : 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 + %c1_f32 = scalar.constant 1.0 : f32 + %cneg1_f32 = scalar.constant -1.0 : f32 + %shifted_word = scalar.shrui %grid_word, %shift : i32 + %value_i32 = scalar.andi %shifted_word, %c255_i32 : i32 + %value_f32 = scalar.uitofp %value_i32 : i32 to f32 + %sign_mask = scalar.andi %signs, %sign_bit : i32 + %negative = scalar.cmpi ne, %sign_mask, %c0_i32 : i32 + %sign = scf.select %negative, %cneg1_f32, %c1_f32 : f32 + %signed = scalar.mulf %value_f32, %sign : f32 + %result = scalar.mulf %signed, %scale : f32 + func.return %result : f32 +} + +func.def inline @ggml_iq2s_f32_vector4(%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 + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %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 + %c8_i32 = scalar.constant 8 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c24_i32 = scalar.constant 24 : i32 + %c255_i32 = scalar.constant 255 : 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> + %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 + %word_index_i32 = index.cast %word_index : 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 = scalar.ori %qs, %qh_masked : i32 + %grid_word = func.call @ggml_iq2s_grid_lookup_i32(%grid_index, %word_index_i32) : (i32, i32) -> (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_f16_vector4(%weight: buffer, %row_byte_base: offset, %iq2_block: index, %iq2_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_iq2s_f32_vector4(%weight, %row_byte_base, %iq2_block, %iq2_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + %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 + %grid_low2 = scalar.constant 261.0 : f32 + %grid_low3 = scalar.constant 267.0 : f32 + %grid_low4 = scalar.constant 271.0 : f32 + %grid_low5 = scalar.constant 769.0 : f32 + %grid_low6 = scalar.constant 771.0 : f32 + %grid_low7 = scalar.constant 773.0 : f32 + %grid_low8 = scalar.constant 777.0 : f32 + %grid_low9 = scalar.constant 781.0 : f32 + %grid_low10 = scalar.constant 1281.0 : f32 + %grid_low11 = scalar.constant 1283.0 : f32 + %grid_low12 = scalar.constant 1291.0 : f32 + %grid_low13 = scalar.constant 1799.0 : f32 + %grid_low14 = scalar.constant 2305.0 : f32 + %grid_low15 = scalar.constant 2309.0 : f32 + %grid_low16 = scalar.constant 2315.0 : f32 + %grid_low17 = scalar.constant 2319.0 : f32 + %grid_low18 = scalar.constant 2819.0 : f32 + %grid_low19 = scalar.constant 2823.0 : f32 + %grid_low20 = scalar.constant 3329.0 : f32 + %grid_low21 = scalar.constant 3333.0 : f32 + %grid_low22 = scalar.constant 3843.0 : f32 + %grid_low23 = scalar.constant 3849.0 : f32 + %grid_low24 = scalar.constant 3855.0 : f32 + %grid_low28 = scalar.constant 265.0 : f32 + %grid_low31 = scalar.constant 779.0 : f32 + %grid_low33 = scalar.constant 1287.0 : f32 + %grid_low34 = scalar.constant 1295.0 : f32 + %grid_low35 = scalar.constant 1795.0 : f32 + %grid_low36 = scalar.constant 1803.0 : f32 + %grid_low37 = scalar.constant 2313.0 : f32 + %grid_low38 = scalar.constant 3331.0 : f32 + %grid_low39 = scalar.constant 3339.0 : f32 + %grid_low40 = scalar.constant 3845.0 : f32 + %grid_low46 = scalar.constant 775.0 : f32 + %grid_low50 = scalar.constant 1793.0 : f32 + %grid_low51 = scalar.constant 1801.0 : f32 + %grid_low57 = scalar.constant 3841.0 : f32 + %grid_low58 = scalar.constant 3847.0 : f32 + %grid_low59 = scalar.constant 263.0 : f32 + %grid_low63 = scalar.constant 1285.0 : f32 + %grid_low66 = scalar.constant 1805.0 : f32 + %grid_low68 = scalar.constant 2817.0 : f32 + %grid_low69 = scalar.constant 2821.0 : f32 + %grid_low70 = scalar.constant 3343.0 : f32 + %grid_low72 = scalar.constant 3851.0 : f32 + %grid_low75 = scalar.constant 783.0 : f32 + %grid_low77 = scalar.constant 1289.0 : f32 + %grid_low78 = scalar.constant 1797.0 : f32 + %grid_low80 = scalar.constant 2311.0 : f32 + %grid_low87 = scalar.constant 1293.0 : f32 + %grid_low89 = scalar.constant 2307.0 : f32 + %grid_low92 = scalar.constant 3341.0 : f32 + %grid_low94 = scalar.constant 269.0 : f32 + %grid_low108 = scalar.constant 2825.0 : f32 + %grid_low124 = scalar.constant 2829.0 : f32 + %grid_low140 = scalar.constant 3853.0 : f32 + %grid_low150 = scalar.constant 2827.0 : f32 + %grid_low172 = scalar.constant 2317.0 : f32 + %grid_low236 = scalar.constant 1807.0 : f32 + %grid_low239 = scalar.constant 2831.0 : f32 + %grid_low294 = scalar.constant 3335.0 : f32 + %grid_low331 = scalar.constant 3337.0 : f32 + %grid_low0_chunk = vector.from_elements %grid_low0, %grid_low1, %grid_low2, %grid_low3, %grid_low4, %grid_low5, %grid_low6, %grid_low7, %grid_low8, %grid_low9, %grid_low10, %grid_low11, %grid_low12, %grid_low13, %grid_low14, %grid_low15, %grid_low16, %grid_low17, %grid_low18, %grid_low19, %grid_low20, %grid_low21, %grid_low22, %grid_low23, %grid_low24, %grid_low0, %grid_low1, %grid_low2, %grid_low28, %grid_low5, %grid_low6, %grid_low31 : vector<32xf32> + %grid_low1_chunk = vector.from_elements %grid_low10, %grid_low33, %grid_low34, %grid_low35, %grid_low36, %grid_low37, %grid_low38, %grid_low39, %grid_low40, %grid_low0, %grid_low1, %grid_low3, %grid_low4, %grid_low5, %grid_low46, %grid_low9, %grid_low11, %grid_low12, %grid_low50, %grid_low51, %grid_low15, %grid_low16, %grid_low17, %grid_low18, %grid_low19, %grid_low57, %grid_low58, %grid_low59, %grid_low6, %grid_low31, %grid_low10, %grid_low63 : vector<32xf32> + %grid_low2_chunk = vector.from_elements %grid_low35, %grid_low13, %grid_low66, %grid_low37, %grid_low68, %grid_low69, %grid_low70, %grid_low22, %grid_low72, %grid_low0, %grid_low46, %grid_low75, %grid_low11, %grid_low77, %grid_low78, %grid_low14, %grid_low80, %grid_low18, %grid_low57, %grid_low2, %grid_low28, %grid_low10, %grid_low63, %grid_low87, %grid_low13, %grid_low89, %grid_low16, %grid_low17, %grid_low92, %grid_low58, %grid_low94, %grid_low6 : vector<32xf32> + %grid_low3_chunk = vector.from_elements %grid_low46, %grid_low35, %grid_low69, %grid_low22, %grid_low0, %grid_low2, %grid_low28, %grid_low10, %grid_low63, %grid_low87, %grid_low13, %grid_low68, %grid_low108, %grid_low0, %grid_low1, %grid_low2, %grid_low28, %grid_low5, %grid_low6, %grid_low46, %grid_low31, %grid_low75, %grid_low10, %grid_low63, %grid_low35, %grid_low51, %grid_low66, %grid_low108, %grid_low124, %grid_low38, %grid_low40, %grid_low0 : vector<32xf32> + %grid_low4_chunk = vector.from_elements %grid_low1, %grid_low59, %grid_low94, %grid_low5, %grid_low8, %grid_low11, %grid_low50, %grid_low13, %grid_low89, %grid_low68, %grid_low69, %grid_low57, %grid_low140, %grid_low0, %grid_low7, %grid_low31, %grid_low75, %grid_low10, %grid_low77, %grid_low78, %grid_low14, %grid_low80, %grid_low150, %grid_low20, %grid_low40, %grid_low1, %grid_low28, %grid_low4, %grid_low5, %grid_low46, %grid_low11, %grid_low34 : vector<32xf32> + %grid_low5_chunk = vector.from_elements %grid_low50, %grid_low51, %grid_low89, %grid_low21, %grid_low57, %grid_low59, %grid_low3, %grid_low7, %grid_low8, %grid_low35, %grid_low13, %grid_low15, %grid_low172, %grid_low68, %grid_low108, %grid_low1, %grid_low5, %grid_low46, %grid_low11, %grid_low50, %grid_low78, %grid_low18, %grid_low10, %grid_low77, %grid_low34, %grid_low37, %grid_low172, %grid_low1, %grid_low59, %grid_low5, %grid_low7, %grid_low11 : vector<32xf32> + %grid_low6_chunk = vector.from_elements %grid_low36, %grid_low89, %grid_low21, %grid_low57, %grid_low0, %grid_low1, %grid_low59, %grid_low3, %grid_low4, %grid_low5, %grid_low7, %grid_low8, %grid_low9, %grid_low11, %grid_low33, %grid_low34, %grid_low50, %grid_low78, %grid_low89, %grid_low80, %grid_low16, %grid_low68, %grid_low69, %grid_low70, %grid_low57, %grid_low58, %grid_low72, %grid_low0, %grid_low2, %grid_low5, %grid_low46, %grid_low75 : vector<32xf32> + %grid_low7_chunk = vector.from_elements %grid_low63, %grid_low12, %grid_low35, %grid_low51, %grid_low15, %grid_low18, %grid_low1, %grid_low28, %grid_low4, %grid_low11, %grid_low33, %grid_low50, %grid_low236, %grid_low89, %grid_low19, %grid_low239, %grid_low22, %grid_low23, %grid_low0, %grid_low2, %grid_low3, %grid_low6, %grid_low63, %grid_low77, %grid_low35, %grid_low13, %grid_low15, %grid_low68, %grid_low92, %grid_low1, %grid_low4, %grid_low10 : vector<32xf32> + %grid_low8_chunk = vector.from_elements %grid_low33, %grid_low78, %grid_low36, %grid_low89, %grid_low40, %grid_low72, %grid_low28, %grid_low6, %grid_low63, %grid_low236, %grid_low14, %grid_low19, %grid_low57, %grid_low0, %grid_low2, %grid_low4, %grid_low11, %grid_low150, %grid_low38, %grid_low3, %grid_low6, %grid_low87, %grid_low50, %grid_low80, %grid_low68, %grid_low2, %grid_low6, %grid_low46, %grid_low31, %grid_low75, %grid_low63, %grid_low35 : vector<32xf32> + %grid_low9_chunk = vector.from_elements %grid_low13, %grid_low36, %grid_low15, %grid_low37, %grid_low17, %grid_low18, %grid_low294, %grid_low22, %grid_low1, %grid_low59, %grid_low3, %grid_low8, %grid_low11, %grid_low33, %grid_low14, %grid_low20, %grid_low40, %grid_low140, %grid_low0, %grid_low7, %grid_low10, %grid_low78, %grid_low51, %grid_low68, %grid_low1, %grid_low5, %grid_low8, %grid_low11, %grid_low33, %grid_low34, %grid_low50, %grid_low89 : vector<32xf32> + %grid_low10_chunk = vector.from_elements %grid_low80, %grid_low17, %grid_low150, %grid_low58, %grid_low59, %grid_low6, %grid_low9, %grid_low63, %grid_low35, %grid_low69, %grid_low20, %grid_low331, %grid_low1, %grid_low5, %grid_low7, %grid_low12, %grid_low78, %grid_low37, %grid_low124, %grid_low58, %grid_low9, %grid_low89, %grid_low1, %grid_low59, %grid_low10, %grid_low63, %grid_low36, %grid_low0, %grid_low28, %grid_low7, %grid_low10, %grid_low77 : vector<32xf32> + %grid_low11_chunk = vector.from_elements %grid_low34, %grid_low78, %grid_low89, %grid_low68, %grid_low57, %grid_low2, %grid_low4, %grid_low6, %grid_low46, %grid_low63, %grid_low50, %grid_low36, %grid_low80, %grid_low18, %grid_low150, %grid_low1, %grid_low59, %grid_low5, %grid_low31, %grid_low11, %grid_low13, %grid_low14, %grid_low239, %grid_low21, %grid_low57, %grid_low28, %grid_low6, %grid_low46, %grid_low10, %grid_low63, %grid_low35, %grid_low36 : vector<32xf32> + %grid_low12_chunk = vector.from_elements %grid_low0, %grid_low2, %grid_low77, %grid_low236, %grid_low14, %grid_low22, %grid_low3, %grid_low4, %grid_low11, %grid_low21, %grid_low46, %grid_low51, %grid_low20, %grid_low5, %grid_low31, %grid_low50, %grid_low80, %grid_low18, %grid_low2, %grid_low5, %grid_low8, %grid_low63, %grid_low14, %grid_low37, %grid_low17, %grid_low69, %grid_low92, %grid_low23, %grid_low1, %grid_low59, %grid_low3, %grid_low7 : vector<32xf32> + %grid_low13_chunk = vector.from_elements %grid_low11, %grid_low78, %grid_low40, %grid_low0, %grid_low6, %grid_low33, %grid_low50, %grid_low66, %grid_low19, %grid_low2, %grid_low4, %grid_low5, %grid_low34, %grid_low37, %grid_low18, %grid_low39, %grid_low58, %grid_low1, %grid_low28, %grid_low10, %grid_low78, %grid_low172, %grid_low7, %grid_low87, %grid_low18, %grid_low19, %grid_low15, %grid_low2, %grid_low28, %grid_low63, %grid_low6, %grid_low46 : vector<32xf32> + %grid_low14_chunk = vector.from_elements %grid_low31, %grid_low35, %grid_low13, %grid_low20, %grid_low0, %grid_low10, %grid_low34, %grid_low331, %grid_low7, %grid_low51, %grid_low15, %grid_low150, %grid_low21, %grid_low57, %grid_low0, %grid_low8, %grid_low11, %grid_low14, %grid_low12, %grid_low80, %grid_low21, %grid_low0, %grid_low59, %grid_low51, %grid_low20, %grid_low3, %grid_low14, %grid_low6, %grid_low46, %grid_low0, %grid_low28, %grid_low4 : vector<32xf32> + %grid_low15_chunk = vector.from_elements %grid_low10, %grid_low63, %grid_low66, %grid_low14, %grid_low108, %grid_low21, %grid_low2, %grid_low6, %grid_low77, %grid_low80, %grid_low16, %grid_low1, %grid_low28, %grid_low5, %grid_low9, %grid_low11, %grid_low50, %grid_low18, %grid_low2, %grid_low78, %grid_low36, %grid_low19, %grid_low1, %grid_low3, %grid_low46, %grid_low10, %grid_low68, %grid_low63, %grid_low15, %grid_low2, %grid_low35, %grid_low0 : vector<32xf32> + %grid_high0_chunk = vector.from_elements %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low0, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1 : vector<32xf32> + %grid_high1_chunk = vector.from_elements %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low1, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low2, %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low59 : vector<32xf32> + %grid_high2_chunk = vector.from_elements %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low59, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low28, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low3, %grid_low94, %grid_low94 : vector<32xf32> + %grid_high3_chunk = vector.from_elements %grid_low94, %grid_low94, %grid_low94, %grid_low94, %grid_low4, %grid_low4, %grid_low4, %grid_low4, %grid_low4, %grid_low4, %grid_low4, %grid_low4, %grid_low4, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low5, %grid_low6 : vector<32xf32> + %grid_high4_chunk = vector.from_elements %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low6, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low7, %grid_low46, %grid_low46, %grid_low46, %grid_low46, %grid_low46, %grid_low46, %grid_low46 : vector<32xf32> + %grid_high5_chunk = vector.from_elements %grid_low46, %grid_low46, %grid_low46, %grid_low46, %grid_low46, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low8, %grid_low31, %grid_low31, %grid_low31, %grid_low31, %grid_low31, %grid_low31, %grid_low31, %grid_low9, %grid_low9, %grid_low9, %grid_low9, %grid_low9, %grid_low75, %grid_low75, %grid_low75, %grid_low75, %grid_low75 : vector<32xf32> + %grid_high6_chunk = vector.from_elements %grid_low75, %grid_low75, %grid_low75, %grid_low75, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low10, %grid_low11, %grid_low11, %grid_low11, %grid_low11, %grid_low11 : vector<32xf32> + %grid_high7_chunk = vector.from_elements %grid_low11, %grid_low11, %grid_low11, %grid_low11, %grid_low11, %grid_low11, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low63, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low33, %grid_low77, %grid_low77, %grid_low77 : vector<32xf32> + %grid_high8_chunk = vector.from_elements %grid_low77, %grid_low77, %grid_low77, %grid_low77, %grid_low77, %grid_low77, %grid_low12, %grid_low12, %grid_low12, %grid_low12, %grid_low12, %grid_low12, %grid_low12, %grid_low87, %grid_low87, %grid_low87, %grid_low87, %grid_low87, %grid_low87, %grid_low34, %grid_low34, %grid_low34, %grid_low34, %grid_low34, %grid_low34, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50 : vector<32xf32> + %grid_high9_chunk = vector.from_elements %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low50, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low35, %grid_low78, %grid_low78, %grid_low78, %grid_low78, %grid_low78, %grid_low78, %grid_low13, %grid_low13, %grid_low13, %grid_low13, %grid_low13, %grid_low13, %grid_low13, %grid_low13 : vector<32xf32> + %grid_high10_chunk = vector.from_elements %grid_low13, %grid_low13, %grid_low13, %grid_low13, %grid_low51, %grid_low51, %grid_low51, %grid_low51, %grid_low51, %grid_low51, %grid_low51, %grid_low51, %grid_low36, %grid_low36, %grid_low36, %grid_low36, %grid_low36, %grid_low36, %grid_low36, %grid_low36, %grid_low66, %grid_low66, %grid_low236, %grid_low236, %grid_low236, %grid_low236, %grid_low236, %grid_low14, %grid_low14, %grid_low14, %grid_low14, %grid_low14 : vector<32xf32> + %grid_high11_chunk = vector.from_elements %grid_low14, %grid_low14, %grid_low14, %grid_low14, %grid_low14, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low89, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low15, %grid_low80, %grid_low80, %grid_low80, %grid_low80, %grid_low80, %grid_low80, %grid_low80 : vector<32xf32> + %grid_high12_chunk = vector.from_elements %grid_low37, %grid_low37, %grid_low37, %grid_low37, %grid_low37, %grid_low37, %grid_low16, %grid_low16, %grid_low16, %grid_low16, %grid_low172, %grid_low172, %grid_low172, %grid_low17, %grid_low17, %grid_low17, %grid_low17, %grid_low17, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low68, %grid_low18, %grid_low18, %grid_low18, %grid_low18 : vector<32xf32> + %grid_high13_chunk = vector.from_elements %grid_low18, %grid_low18, %grid_low18, %grid_low69, %grid_low69, %grid_low69, %grid_low69, %grid_low69, %grid_low69, %grid_low19, %grid_low19, %grid_low19, %grid_low19, %grid_low19, %grid_low19, %grid_low19, %grid_low19, %grid_low108, %grid_low108, %grid_low108, %grid_low108, %grid_low108, %grid_low150, %grid_low150, %grid_low150, %grid_low150, %grid_low124, %grid_low239, %grid_low239, %grid_low239, %grid_low20, %grid_low20 : vector<32xf32> + %grid_high14_chunk = vector.from_elements %grid_low20, %grid_low20, %grid_low20, %grid_low20, %grid_low38, %grid_low38, %grid_low38, %grid_low38, %grid_low21, %grid_low21, %grid_low21, %grid_low21, %grid_low21, %grid_low21, %grid_low294, %grid_low294, %grid_low294, %grid_low294, %grid_low331, %grid_low331, %grid_low331, %grid_low39, %grid_low39, %grid_low39, %grid_low39, %grid_low92, %grid_low92, %grid_low70, %grid_low70, %grid_low57, %grid_low57, %grid_low57 : vector<32xf32> + %grid_high15_chunk = vector.from_elements %grid_low57, %grid_low57, %grid_low57, %grid_low57, %grid_low57, %grid_low57, %grid_low22, %grid_low22, %grid_low22, %grid_low22, %grid_low22, %grid_low40, %grid_low40, %grid_low40, %grid_low40, %grid_low40, %grid_low40, %grid_low40, %grid_low58, %grid_low58, %grid_low58, %grid_low58, %grid_low23, %grid_low23, %grid_low23, %grid_low23, %grid_low23, %grid_low72, %grid_low72, %grid_low140, %grid_low140, %grid_low24 : vector<32xf32> + func.return %grid_low0_chunk, %grid_low1_chunk, %grid_low2_chunk, %grid_low3_chunk, %grid_low4_chunk, %grid_low5_chunk, %grid_low6_chunk, %grid_low7_chunk, %grid_low8_chunk, %grid_low9_chunk, %grid_low10_chunk, %grid_low11_chunk, %grid_low12_chunk, %grid_low13_chunk, %grid_low14_chunk, %grid_low15_chunk, %grid_high0_chunk, %grid_high1_chunk, %grid_high2_chunk, %grid_high3_chunk, %grid_high4_chunk, %grid_high5_chunk, %grid_high6_chunk, %grid_high7_chunk, %grid_high8_chunk, %grid_high9_chunk, %grid_high10_chunk, %grid_high11_chunk, %grid_high12_chunk, %grid_high13_chunk, %grid_high14_chunk, %grid_high15_chunk : vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32> +} + +func.def inline @ggml_iq3s_grid_lookup_i32(%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>, %grid_index: i32) -> (i32) { + %c5_i32 = scalar.constant 5 : i32 + %c31_i32 = scalar.constant 31 : 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 + %c9_i32 = scalar.constant 9 : i32 + %c10_i32 = scalar.constant 10 : i32 + %c11_i32 = scalar.constant 11 : i32 + %c12_i32 = scalar.constant 12 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c14_i32 = scalar.constant 14 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %chunk_i32 = scalar.shrui %grid_index, %c5_i32 : i32 + %lane_i32 = scalar.andi %grid_index, %c31_i32 : i32 + %codes = vector.from_elements %lane_i32 : vector<1xi32> + %is_chunk1 = scalar.cmpi eq, %chunk_i32, %c1_i32 : i32 + %is_chunk2 = scalar.cmpi eq, %chunk_i32, %c2_i32 : i32 + %is_chunk3 = scalar.cmpi eq, %chunk_i32, %c3_i32 : i32 + %is_chunk4 = scalar.cmpi eq, %chunk_i32, %c4_i32 : i32 + %is_chunk5 = scalar.cmpi eq, %chunk_i32, %c5_i32 : i32 + %is_chunk6 = scalar.cmpi eq, %chunk_i32, %c6_i32 : i32 + %is_chunk7 = scalar.cmpi eq, %chunk_i32, %c7_i32 : i32 + %is_chunk8 = scalar.cmpi eq, %chunk_i32, %c8_i32 : i32 + %is_chunk9 = scalar.cmpi eq, %chunk_i32, %c9_i32 : i32 + %is_chunk10 = scalar.cmpi eq, %chunk_i32, %c10_i32 : i32 + %is_chunk11 = scalar.cmpi eq, %chunk_i32, %c11_i32 : i32 + %is_chunk12 = scalar.cmpi eq, %chunk_i32, %c12_i32 : i32 + %is_chunk13 = scalar.cmpi eq, %chunk_i32, %c13_i32 : i32 + %is_chunk14 = scalar.cmpi eq, %chunk_i32, %c14_i32 : i32 + %is_chunk15 = scalar.cmpi eq, %chunk_i32, %c15_i32 : i32 + %selected_low1 = scf.select %is_chunk1, %iq3s_grid_low1, %iq3s_grid_low0 : vector<32xf32> + %selected_low2 = scf.select %is_chunk2, %iq3s_grid_low2, %selected_low1 : vector<32xf32> + %selected_low3 = scf.select %is_chunk3, %iq3s_grid_low3, %selected_low2 : vector<32xf32> + %selected_low4 = scf.select %is_chunk4, %iq3s_grid_low4, %selected_low3 : vector<32xf32> + %selected_low5 = scf.select %is_chunk5, %iq3s_grid_low5, %selected_low4 : vector<32xf32> + %selected_low6 = scf.select %is_chunk6, %iq3s_grid_low6, %selected_low5 : vector<32xf32> + %selected_low7 = scf.select %is_chunk7, %iq3s_grid_low7, %selected_low6 : vector<32xf32> + %selected_low8 = scf.select %is_chunk8, %iq3s_grid_low8, %selected_low7 : vector<32xf32> + %selected_low9 = scf.select %is_chunk9, %iq3s_grid_low9, %selected_low8 : vector<32xf32> + %selected_low10 = scf.select %is_chunk10, %iq3s_grid_low10, %selected_low9 : vector<32xf32> + %selected_low11 = scf.select %is_chunk11, %iq3s_grid_low11, %selected_low10 : vector<32xf32> + %selected_low12 = scf.select %is_chunk12, %iq3s_grid_low12, %selected_low11 : vector<32xf32> + %selected_low13 = scf.select %is_chunk13, %iq3s_grid_low13, %selected_low12 : vector<32xf32> + %selected_low14 = scf.select %is_chunk14, %iq3s_grid_low14, %selected_low13 : vector<32xf32> + %selected_low = scf.select %is_chunk15, %iq3s_grid_low15, %selected_low14 : vector<32xf32> + %selected_high1 = scf.select %is_chunk1, %iq3s_grid_high1, %iq3s_grid_high0 : vector<32xf32> + %selected_high2 = scf.select %is_chunk2, %iq3s_grid_high2, %selected_high1 : vector<32xf32> + %selected_high3 = scf.select %is_chunk3, %iq3s_grid_high3, %selected_high2 : vector<32xf32> + %selected_high4 = scf.select %is_chunk4, %iq3s_grid_high4, %selected_high3 : vector<32xf32> + %selected_high5 = scf.select %is_chunk5, %iq3s_grid_high5, %selected_high4 : vector<32xf32> + %selected_high6 = scf.select %is_chunk6, %iq3s_grid_high6, %selected_high5 : vector<32xf32> + %selected_high7 = scf.select %is_chunk7, %iq3s_grid_high7, %selected_high6 : vector<32xf32> + %selected_high8 = scf.select %is_chunk8, %iq3s_grid_high8, %selected_high7 : vector<32xf32> + %selected_high9 = scf.select %is_chunk9, %iq3s_grid_high9, %selected_high8 : vector<32xf32> + %selected_high10 = scf.select %is_chunk10, %iq3s_grid_high10, %selected_high9 : vector<32xf32> + %selected_high11 = scf.select %is_chunk11, %iq3s_grid_high11, %selected_high10 : vector<32xf32> + %selected_high12 = scf.select %is_chunk12, %iq3s_grid_high12, %selected_high11 : vector<32xf32> + %selected_high13 = scf.select %is_chunk13, %iq3s_grid_high13, %selected_high12 : vector<32xf32> + %selected_high14 = scf.select %is_chunk14, %iq3s_grid_high14, %selected_high13 : vector<32xf32> + %selected_high = scf.select %is_chunk15, %iq3s_grid_high15, %selected_high14 : vector<32xf32> + %low_values = vector.table.lookup %selected_low[%codes] : vector<32xf32>, vector<1xi32> -> vector<1xf32> + %high_values = vector.table.lookup %selected_high[%codes] : vector<32xf32>, vector<1xi32> -> vector<1xf32> + %low_f32 = vector.extract %low_values[0] : vector<1xf32> -> f32 + %high_f32 = vector.extract %high_values[0] : vector<1xf32> -> f32 + %low = scalar.fptoui %low_f32 : f32 to i32 + %high = scalar.fptoui %high_f32 : f32 to i32 + %high_shifted = scalar.shli %high, %c16_i32 : i32 + %word = scalar.ori %low, %high_shifted : i32 + func.return %word : i32 +} + +func.def inline @ggml_iq3s_signed_value_f32(%grid_word: i32, %shift: i32, %shifted_signs: i32, %sign_bit: i32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c255_i32 = scalar.constant 255 : i32 + %shifted_word = scalar.shrui %grid_word, %shift : i32 + %value_i32 = scalar.andi %shifted_word, %c255_i32 : i32 + %negative_value = scalar.subi %c0_i32, %value_i32 : i32 + %masked_sign = scalar.andi %shifted_signs, %sign_bit : i32 + %is_negative = scalar.cmpi ne, %masked_sign, %c0_i32 : i32 + %signed_value = scf.select %is_negative, %negative_value, %value_i32 : i32 + %value = scalar.sitofp %signed_value : i32 to f32 + func.return %value : f32 +} + +func.def inline @ggml_iq3s_f32_vector4(%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: buffer, %row_byte_base: offset, %iq3_block: index, %iq3_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c0_i32 = scalar.constant 0 : i32 + %c2 = index.constant 2 : index + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c4 = index.constant 4 : index + %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> + %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 = scalar.ori %q_i32, %qh_high : i32 + %grid_word = func.call @ggml_iq3s_grid_lookup_i32(%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, %grid_index) : (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, i32) -> (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_iq3s_f16_vector4(%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: buffer, %row_byte_base: offset, %iq3_block: index, %iq3_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_iq3s_f32_vector4(%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, %row_byte_base, %iq3_block, %iq3_group, %packet) : (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, 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_iq4nl_table_i8() -> (vector<16xi8>) { + %v0 = scalar.constant -127 : i8 + %v1 = scalar.constant -104 : i8 + %v2 = scalar.constant -83 : i8 + %v3 = scalar.constant -65 : i8 + %v4 = scalar.constant -49 : i8 + %v5 = scalar.constant -35 : i8 + %v6 = scalar.constant -22 : i8 + %v7 = scalar.constant -10 : i8 + %v8 = scalar.constant 1 : i8 + %v9 = scalar.constant 13 : i8 + %v10 = scalar.constant 25 : i8 + %v11 = scalar.constant 38 : i8 + %v12 = scalar.constant 53 : i8 + %v13 = scalar.constant 69 : i8 + %v14 = scalar.constant 89 : i8 + %v15 = scalar.constant 113 : i8 + %value_table = vector.from_elements %v0, %v1, %v2, %v3, %v4, %v5, %v6, %v7, %v8, %v9, %v10, %v11, %v12, %v13, %v14, %v15 : vector<16xi8> + func.return %value_table : vector<16xi8> +} + +func.def inline @ggml_iq4nl_lookup_i8_vector4(%value_table: vector<16xi8>, %codes: vector<4xi8>) -> (vector<4xi8>) { + %values = vector.table.lookup %value_table[%codes] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + func.return %values : vector<4xi8> +} + +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 + %q_word_vector = vector.bitcast %q_bytes : vector<4xi8> to vector<1xi32> + %word = vector.extract %q_word_vector[0] : vector<1xi32> -> i32 + %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> +} + +func.def inline @ggml_iq4nl_f32_vector4(%iq4nl_table: vector<16xi8>, %weight: buffer, %row_byte_base: offset, %iq4_block: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 18 : offset + %code_offset = index.constant 2 : offset + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %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 + %uses_high = index.cmp uge, %bounded_packet, %c4 : index + %q_bytes = vector.load %code_view[%code_index] : view<16xi8> -> vector<4xi8> + %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_vector = vector.splat %d : vector<4xf32> + %result = vector.mulf %d_vector, %lookup : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_iq4nl_f16_vector4(%iq4nl_table: vector<16xi8>, %weight: buffer, %row_byte_base: offset, %iq4_block: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_iq4nl_f32_vector4(%iq4nl_table, %weight, %row_byte_base, %iq4_block, %packet) : (vector<16xi8>, buffer, offset, index, index) -> (vector<4xf32>) + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +// Packs either nibble from four adjacent GGUF IQ4_XS bytes for the gfx11 table lookup. +// Preserve GGUF code indices for the logical table lookup. +func.def inline @ggml_iq4xs_table_codes4(%code_word: vector<1xi32>, %uses_high: i1) -> (vector<4xi8>) { + %c4_i32 = scalar.constant 4 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c12_i32 = scalar.constant 12 : i32 + %c16_i32 = scalar.constant 16 : 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 + %word = vector.extract %code_word[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> + func.return %codes : vector<4xi8> +} + +func.def inline @ggml_iq4xs_f32_vector4(%iq4nl_table: vector<16xi8>, %weight: buffer, %row_byte_base: offset, %iq4_block: index, %iq4_group: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 136 : offset + %scales_l_offset = index.constant 4 : offset + %code_offset = index.constant 8 : offset + %c3_i32 = scalar.constant 3 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c32_i32 = scalar.constant 32 : i32 + %c65535_i32 = scalar.constant 65535 : i32 + %bounded_group = index.assume %iq4_group [range(%iq4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scales_l_byte_base = index.add %block_byte_base, %scales_l_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xi32> + %scales_l_view = buffer.view %weight[%scales_l_byte_base] : buffer -> view<1xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header = view.load %header_view[%c0] : view<1xi32> -> i32 + %header_vector = vector.from_elements %header : vector<1xi32> + %header_halves = vector.bitcast %header_vector : vector<1xi32> to vector<2xf16> + %d_f16 = vector.extract %header_halves[0] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %scales_l = view.load %scales_l_view[%c0] : view<1xi32> -> i32 + %scales_h_shifted = scalar.shrui %header, %c16_i32 : i32 + %scales_h = scalar.andi %scales_h_shifted, %c65535_i32 : i32 + %scale_l_shift_index = index.mul %bounded_group, %c4 : index + %scale_l_shift = index.cast %scale_l_shift_index : index to i32 + %scale_l_shifted = scalar.shrui %scales_l, %scale_l_shift : i32 + %scale_l = scalar.andi %scale_l_shifted, %c15_i32 : i32 + %scale_h_shift_index = index.mul %bounded_group, %c2 : index + %scale_h_shift = index.cast %scale_h_shift_index : index to i32 + %scale_h_shifted = scalar.shrui %scales_h, %scale_h_shift : i32 + %scale_h = scalar.andi %scale_h_shifted, %c3_i32 : i32 + %scale_h_high = scalar.shli %scale_h, %c4_i32 : i32 + %local_scale = scalar.ori %scale_l, %scale_h_high : i32 + %centered_scale = scalar.subi %local_scale, %c32_i32 : i32 + %centered_scale_f32 = scalar.sitofp %centered_scale : i32 to f32 + %block_scale = scalar.mulf %d, %centered_scale_f32 : f32 + %packet_page = index.rem %bounded_packet, %c4 : index + %q_word_base = index.mul %bounded_group, %c4 : index + %q_word_index0 = index.add %q_word_base, %packet_page : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %uses_high = index.cmp uge, %bounded_packet, %c4 : index + %codes = func.call @ggml_iq4xs_table_codes4(%q_word, %uses_high) : (vector<1xi32>, i1) -> (vector<4xi8>) + %lookup_i8 = func.call @ggml_iq4nl_lookup_i8_vector4(%iq4nl_table, %codes) : (vector<16xi8>, vector<4xi8>) -> (vector<4xi8>) + %lookup0_i8 = vector.extract %lookup_i8[0] : vector<4xi8> -> i8 + %lookup1_i8 = vector.extract %lookup_i8[1] : vector<4xi8> -> i8 + %lookup2_i8 = vector.extract %lookup_i8[2] : vector<4xi8> -> i8 + %lookup3_i8 = vector.extract %lookup_i8[3] : vector<4xi8> -> i8 + %lookup0_i32 = scalar.extsi %lookup0_i8 : i8 to i32 + %lookup1_i32 = scalar.extsi %lookup1_i8 : i8 to i32 + %lookup2_i32 = scalar.extsi %lookup2_i8 : i8 to i32 + %lookup3_i32 = scalar.extsi %lookup3_i8 : i8 to i32 + %lookup0 = scalar.sitofp %lookup0_i32 : i32 to f32 + %lookup1 = scalar.sitofp %lookup1_i32 : i32 to f32 + %lookup2 = scalar.sitofp %lookup2_i32 : i32 to f32 + %lookup3 = scalar.sitofp %lookup3_i32 : i32 to f32 + %value0 = scalar.mulf %block_scale, %lookup0 : f32 + %value1 = scalar.mulf %block_scale, %lookup1 : f32 + %value2 = scalar.mulf %block_scale, %lookup2 : f32 + %value3 = scalar.mulf %block_scale, %lookup3 : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +func.def inline @ggml_iq4xs_f16_vector4(%iq4nl_table: vector<16xi8>, %weight: buffer, %row_byte_base: offset, %iq4_block: index, %iq4_group: index, %packet: index) -> (vector<4xf16>) { + %values_f32 = func.call @ggml_iq4xs_f32_vector4(%iq4nl_table, %weight, %row_byte_base, %iq4_block, %iq4_group, %packet) : (vector<16xi8>, 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_q6k_f32_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %ql_half_word_base = index.mul %half, %c16 : index + %ql_side_word_add = index.mul %ql_side, %c8 : index + %ql_word_base = index.add %ql_half_word_base, %ql_side_word_add : index + %ql_word_index = index.add %ql_word_base, %bounded_packet : index + %qh_half_word_base = index.mul %half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index = index.add %scale_group_base, %scale_packet_half : index + %ql_word = vector.load %ql_view[%ql_word_index] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + func.return %values : vector<4xf32> +} + +func.def inline @ggml_q8_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 34 : offset + %code_offset = index.constant 2 : offset + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %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<32xi8> + %packet_base = index.mul %bounded_packet, %c4 : index + %codes = vector.load %code_view[%packet_base] : view<32xi8> -> vector<4xi8> + %codes_f32 = vector.sitofp %codes : 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> + %values = vector.mulf %codes_f32, %d_vector : vector<4xf32> + func.return %values : vector<4xf32> +} + +func.def inline @ggml_q8_1_f32_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf32>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 36 : offset + %code_offset = index.constant 4 : offset + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %ds_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi8> + %packet_base = index.mul %bounded_packet, %c4 : index + %codes = vector.load %code_view[%packet_base] : view<32xi8> -> vector<4xi8> + %codes_f32 = vector.sitofp %codes : vector<4xi8> to vector<4xf32> + %d_f16 = view.load %ds_view[%c0] : view<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %d_vector = vector.splat %d : vector<4xf32> + %values = vector.mulf %codes_f32, %d_vector : vector<4xf32> + func.return %values : vector<4xf32> +} + +func.def inline @ggml_f16_f32_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf32>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 4, 1073741824), mul(%input_size, 4)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xf16> + %values_f16 = vector.load %row_view[%k] : view<[%bounded_input_size]xf16> -> vector<4xf16> + %values = vector.extf %values_f16 : vector<4xf16> to vector<4xf32> + func.return %values : vector<4xf32> +} + +func.def inline @ggml_bf16_f32_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf32>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 4, 1073741824), mul(%input_size, 4)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xbf16> + %values_bf16 = vector.load %row_view[%k] : view<[%bounded_input_size]xbf16> -> vector<4xbf16> + %values = vector.extf %values_bf16 : vector<4xbf16> to vector<4xf32> + func.return %values : vector<4xf32> +} + +func.def inline @ggml_f32_f32_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf32>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 4, 1073741824), mul(%input_size, 4)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xf32> + %values = vector.load %row_view[%k] : view<[%bounded_input_size]xf32> -> vector<4xf32> + func.return %values : vector<4xf32> +} + +func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) { + %q1_0_format = index.constant 10 : index + %iq1_s_format = index.constant 19 : index + %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 + %q4_0_format = index.constant 40 : index + %q4_1_format = index.constant 41 : index + %q5_0_format = index.constant 50 : index + %q5_1_format = index.constant 51 : index + %q8_0_format = index.constant 80 : index + %q8_1_format = index.constant 81 : index + %zero_bytes = index.constant 0 : offset + %q1_0_tile_bytes = index.constant 36 : offset + %q3_tile_bytes = index.constant 110 : offset + %q4_tile_bytes = index.constant 144 : offset + %q5_tile_bytes = index.constant 176 : offset + %q6_tile_bytes = index.constant 210 : offset + %q4_0_tile_bytes = index.constant 144 : offset + %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 + %bf16_tile_bytes = index.constant 512 : offset + %f32_tile_bytes = index.constant 1024 : offset + %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index + %is_q3 = index.cmp eq, %weight_format, %c3 : index + %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q5 = index.cmp eq, %weight_format, %c5 : index + %is_q6 = index.cmp eq, %weight_format, %c6 : index + %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index + %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 + %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index + %is_f32 = index.cmp eq, %weight_format, %c32 : index + %selected_q1_0 = scf.select %is_q1_0, %q1_0_tile_bytes, %zero_bytes : offset + %selected_q3 = scf.select %is_q3, %q3_tile_bytes, %selected_q1_0 : offset + %selected_q4 = scf.select %is_q4, %q4_tile_bytes, %selected_q3 : offset + %selected_q5 = scf.select %is_q5, %q5_tile_bytes, %selected_q4 : offset + %selected_q6 = scf.select %is_q6, %q6_tile_bytes, %selected_q5 : offset + %selected_q4_0 = scf.select %is_q4_0, %q4_0_tile_bytes, %selected_q6 : offset + %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_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_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 + %selected = scf.select %is_f32, %f32_tile_bytes, %selected_bf16 : offset + func.return %selected : 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 + %q4_0_format = index.constant 40 : index + %q4_1_format = index.constant 41 : index + %q5_0_format = index.constant 50 : index + %q5_1_format = index.constant 51 : index + %q8_0_format = index.constant 80 : index + %q8_1_format = index.constant 81 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %zero_bytes = index.constant 0 : offset + %q1_0_block_bytes = index.constant 18 : offset + %q4_0_block_bytes = index.constant 18 : offset + %q4_1_block_bytes = index.constant 20 : offset + %q5_0_block_bytes = index.constant 22 : offset + %q5_1_block_bytes = index.constant 24 : offset + %iq4_nl_block_bytes = index.constant 18 : offset + %q8_0_block_bytes = index.constant 34 : offset + %q8_1_block_bytes = index.constant 36 : offset + %f16_element_bytes = index.constant 2 : offset + %f32_element_bytes = index.constant 4 : offset + %uses_k_tile0 = index.cmp eq, %weight_format, %c3 : index + %uses_k_tile1 = index.cmp eq, %weight_format, %c4 : index + %uses_k_tile2 = index.cmp eq, %weight_format, %c5 : index + %uses_k_tile3 = index.cmp eq, %weight_format, %c6 : index + %uses_k_tile4 = index.cmp eq, %weight_format, %c20 : index + %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_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 + %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index + %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index + %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_iq4_nl = index.cmp eq, %weight_format, %c20 : index + %is_q8_0 = index.cmp eq, %weight_format, %q8_0_format : index + %is_q8_1 = index.cmp eq, %weight_format, %q8_1_format : index + %is_f16 = index.cmp eq, %weight_format, %c16 : index + %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index + %is_f32 = index.cmp eq, %weight_format, %c32_format : index + %q1_0_block_count = index.div %hidden_size, %c128 : index + %q1_0_row_bytes = index.scale %q1_0_block_count, %q1_0_block_bytes : index, offset -> offset + %legacy_block_count = index.div %hidden_size, %c32 : index + %q4_0_row_bytes = index.scale %legacy_block_count, %q4_0_block_bytes : index, offset -> offset + %q4_1_row_bytes = index.scale %legacy_block_count, %q4_1_block_bytes : index, offset -> offset + %q5_0_row_bytes = index.scale %legacy_block_count, %q5_0_block_bytes : index, offset -> offset + %q5_1_row_bytes = index.scale %legacy_block_count, %q5_1_block_bytes : index, offset -> offset + %iq4_nl_row_bytes = index.scale %legacy_block_count, %iq4_nl_block_bytes : index, offset -> offset + %q8_0_row_bytes = index.scale %legacy_block_count, %q8_0_block_bytes : index, offset -> offset + %q8_1_row_bytes = index.scale %legacy_block_count, %q8_1_block_bytes : index, offset -> offset + %f16_row_bytes = index.scale %hidden_size, %f16_element_bytes : index, offset -> offset + %f32_row_bytes = index.scale %hidden_size, %f32_element_bytes : index, offset -> offset + %selected_tile = scf.select %uses_k_tile, %tile_row_bytes, %zero_bytes : offset + %selected_q1_0 = scf.select %is_q1_0, %q1_0_row_bytes, %selected_tile : offset + %selected_q4_0 = scf.select %is_q4_0, %q4_0_row_bytes, %selected_q1_0 : offset + %selected_q4_1 = scf.select %is_q4_1, %q4_1_row_bytes, %selected_q4_0 : offset + %selected_q5_0 = scf.select %is_q5_0, %q5_0_row_bytes, %selected_q4_1 : offset + %selected_q5_1 = scf.select %is_q5_1, %q5_1_row_bytes, %selected_q5_0 : offset + %selected_iq4_nl = scf.select %is_iq4_nl, %iq4_nl_row_bytes, %selected_q5_1 : offset + %selected_q8_0 = scf.select %is_q8_0, %q8_0_row_bytes, %selected_iq4_nl : offset + %selected_q8_1 = scf.select %is_q8_1, %q8_1_row_bytes, %selected_q8_0 : offset + %selected_f16 = scf.select %is_f16, %f16_row_bytes, %selected_q8_1 : offset + %selected_bf16 = scf.select %is_bf16, %f16_row_bytes, %selected_f16 : offset + %selected = scf.select %is_f32, %f32_row_bytes, %selected_bf16 : offset + func.return %selected : offset +} + +// Centralizes row-format decoding for FP16 WMMA staging. Each format path is +// independently guarded so inactive layouts do not issue memory loads. +// +// Format ids: +// 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/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>) { + %q1_0_format = index.constant 10 : index + %c3 = index.constant 11 : index + %c4 = index.constant 4 : index + %c5 = index.constant 5 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c20 = index.constant 20 : index + %c21 = index.constant 21 : index + %c22 = index.constant 22 : index + %c23 = index.constant 23 : index + %c16 = index.constant 16 : index + %bf16_format = index.constant 30 : index + %c32 = index.constant 32 : index + %q4_0_format = index.constant 40 : index + %q4_1_format = index.constant 41 : index + %q5_0_format = index.constant 50 : index + %q5_1_format = index.constant 51 : index + %q8_0_format = index.constant 80 : index + %q8_1_format = index.constant 81 : index + %zero = vector.constant 0.0 : vector<4xf16> + // Decode the compact integer format id into one predicate per layout. + %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index + %is_q3 = index.cmp eq, %weight_format, %c3 : index + %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q5 = index.cmp eq, %weight_format, %c5 : index + %is_q6 = index.cmp eq, %weight_format, %c6 : index + %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index + %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_iq2_s = index.cmp eq, %weight_format, %c22 : 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_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 + %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index + %is_f32 = index.cmp eq, %weight_format, %c32 : index + // Shared block index used by 32-value formats that pack eight groups per + // 256-value tile. + %q8_block_base = index.mul %quant_block, %c8 : index + %q8_block = index.add %q8_block_base, %quant_group : index + // Legacy 1-bit format. + %q1_0_values = scf.if %is_q1_0 -> (vector<4xf16>) { + %values = func.call @ggml_q1_0_f16_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + // K-quant family. + %q3_values = scf.if %is_q3 -> (vector<4xf16>) { + %values = func.call @ggml_q3k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q4_values = scf.if %is_q4 -> (vector<4xf16>) { + %values = func.call @ggml_q4k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q6_values = scf.if %is_q6 -> (vector<4xf16>) { + %values = func.call @ggml_q6k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q5_values = scf.if %is_q5 -> (vector<4xf16>) { + %values = func.call @ggml_q5k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + // Legacy 32-value block family. + %q4_0_values = scf.if %is_q4_0 -> (vector<4xf16>) { + %values = func.call @ggml_q4_0_f16_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q4_1_values = scf.if %is_q4_1 -> (vector<4xf16>) { + %values = func.call @ggml_q4_1_f16_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q5_0_values = scf.if %is_q5_0 -> (vector<4xf16>) { + %values = func.call @ggml_q5_0_f16_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q5_1_values = scf.if %is_q5_1 -> (vector<4xf16>) { + %values = func.call @ggml_q5_1_f16_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + // IQ family. + %iq2_s_values = scf.if %is_iq2_s -> (vector<4xf16>) { + %values = func.call @ggml_iq2s_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %iq4_nl_values = scf.if %is_iq4_nl -> (vector<4xf16>) { + %values = func.call @ggml_iq4nl_f16_vector4(%iq4nl_table, %weight, %row_byte_base, %q8_block, %packet) : (vector<16xi8>, buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %iq3_s_values = scf.if %is_iq3_s -> (vector<4xf16>) { + %values = func.call @ggml_iq3s_f16_vector4(%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, %row_byte_base, %quant_block, %quant_group, %packet) : (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %iq4_xs_values = scf.if %is_iq4_xs -> (vector<4xf16>) { + %values = func.call @ggml_iq4xs_f16_vector4(%iq4nl_table, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (vector<16xi8>, buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + // Q8 row formats. + %q8_0_values = scf.if %is_q8_0 -> (vector<4xf16>) { + %values = func.call @ggml_q8_0_f16_vector4(%weight, %row_byte_base, %q8_block, %packet) : (buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %q8_1_values = scf.if %is_q8_1 -> (vector<4xf16>) { + %values = func.call @ggml_q8_1_f16_vector4(%weight, %row_byte_base, %q8_block, %packet) : (buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + // Dense element formats. + %f16_values = scf.if %is_f16 -> (vector<4xf16>) { + %values = func.call @ggml_f16_f16_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %bf16_values = scf.if %is_bf16 -> (vector<4xf16>) { + %values = func.call @ggml_bf16_f16_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + %f32_values = scf.if %is_f32 -> (vector<4xf16>) { + %values = func.call @ggml_f32_f16_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } else { + scf.yield %zero : vector<4xf16> + } + // Select the active decoded vector. Exactly one predicate should be true. + %selected_q1_0 = scf.select %is_q1_0, %q1_0_values, %zero : vector<4xf16> + %selected_q3 = scf.select %is_q3, %q3_values, %selected_q1_0 : vector<4xf16> + %selected_q4 = scf.select %is_q4, %q4_values, %selected_q3 : vector<4xf16> + %selected_q5 = scf.select %is_q5, %q5_values, %selected_q4 : vector<4xf16> + %selected_q6 = scf.select %is_q6, %q6_values, %selected_q5 : vector<4xf16> + %selected_q4_0 = scf.select %is_q4_0, %q4_0_values, %selected_q6 : vector<4xf16> + %selected_q4_1 = scf.select %is_q4_1, %q4_1_values, %selected_q4_0 : vector<4xf16> + %selected_q5_0 = scf.select %is_q5_0, %q5_0_values, %selected_q4_1 : vector<4xf16> + %selected_q5_1 = scf.select %is_q5_1, %q5_1_values, %selected_q5_0 : vector<4xf16> + %selected_iq2_s = scf.select %is_iq2_s, %iq2_s_values, %selected_q5_1 : vector<4xf16> + %selected_iq4_nl = scf.select %is_iq4_nl, %iq4_nl_values, %selected_iq2_s : vector<4xf16> + %selected_iq3_s = scf.select %is_iq3_s, %iq3_s_values, %selected_iq4_nl : vector<4xf16> + %selected_iq4_xs = scf.select %is_iq4_xs, %iq4_xs_values, %selected_iq3_s : vector<4xf16> + %selected_q8_0 = scf.select %is_q8_0, %q8_0_values, %selected_iq4_xs : vector<4xf16> + %selected_q8_1 = scf.select %is_q8_1, %q8_1_values, %selected_q8_0 : vector<4xf16> + %selected_f16 = scf.select %is_f16, %f16_values, %selected_q8_1 : vector<4xf16> + %selected_bf16 = scf.select %is_bf16, %bf16_values, %selected_f16 : vector<4xf16> + %selected = scf.select %is_f32, %f32_values, %selected_bf16 : vector<4xf16> + func.return %selected : vector<4xf16> +} + +// Centralizes row-format decoding for F32 row-gather outputs. +func.def inline @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>) { + %q1_0_format = index.constant 10 : index + %c3 = index.constant 11 : index + %c4 = index.constant 4 : index + %c5 = index.constant 5 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c20 = index.constant 20 : index + %c21 = index.constant 21 : index + %c22 = index.constant 22 : index + %c23 = index.constant 23 : index + %c16 = index.constant 16 : index + %bf16_format = index.constant 30 : index + %c32 = index.constant 32 : index + %q4_0_format = index.constant 40 : index + %q4_1_format = index.constant 41 : index + %q5_0_format = index.constant 50 : index + %q5_1_format = index.constant 51 : index + %q8_0_format = index.constant 80 : index + %q8_1_format = index.constant 81 : index + %zero = vector.constant 0.0 : vector<4xf32> + %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index + %is_q3 = index.cmp eq, %weight_format, %c3 : index + %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q5 = index.cmp eq, %weight_format, %c5 : index + %is_q6 = index.cmp eq, %weight_format, %c6 : index + %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index + %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_iq2_s = index.cmp eq, %weight_format, %c22 : 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_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 + %is_bf16 = index.cmp eq, %weight_format, %bf16_format : index + %is_f32 = index.cmp eq, %weight_format, %c32 : index + %q8_block_base = index.mul %quant_block, %c8 : index + %q8_block = index.add %q8_block_base, %quant_group : index + %q1_0_values = scf.if %is_q1_0 -> (vector<4xf32>) { + %values = func.call @ggml_q1_0_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q3_values = scf.if %is_q3 -> (vector<4xf32>) { + %values = func.call @ggml_q3k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q4_values = scf.if %is_q4 -> (vector<4xf32>) { + %values = func.call @ggml_q4k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q6_values = scf.if %is_q6 -> (vector<4xf32>) { + %values = func.call @ggml_q6k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q5_values = scf.if %is_q5 -> (vector<4xf32>) { + %values = func.call @ggml_q5k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q4_0_values = scf.if %is_q4_0 -> (vector<4xf32>) { + %values = func.call @ggml_q4_0_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q4_1_values = scf.if %is_q4_1 -> (vector<4xf32>) { + %values = func.call @ggml_q4_1_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q5_0_values = scf.if %is_q5_0 -> (vector<4xf32>) { + %values = func.call @ggml_q5_0_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q5_1_values = scf.if %is_q5_1 -> (vector<4xf32>) { + %values = func.call @ggml_q5_1_f32_vector4(%weight, %row_byte_base, %k) : (buffer, offset, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %iq2_s_values = scf.if %is_iq2_s -> (vector<4xf32>) { + %values = func.call @ggml_iq2s_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %iq4_nl_values = scf.if %is_iq4_nl -> (vector<4xf32>) { + %values = func.call @ggml_iq4nl_f32_vector4(%iq4nl_table, %weight, %row_byte_base, %q8_block, %packet) : (vector<16xi8>, buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %iq3_s_values = scf.if %is_iq3_s -> (vector<4xf32>) { + %values = func.call @ggml_iq3s_f32_vector4(%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, %row_byte_base, %quant_block, %quant_group, %packet) : (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %iq4_xs_values = scf.if %is_iq4_xs -> (vector<4xf32>) { + %values = func.call @ggml_iq4xs_f32_vector4(%iq4nl_table, %weight, %row_byte_base, %quant_block, %quant_group, %packet) : (vector<16xi8>, buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q8_0_values = scf.if %is_q8_0 -> (vector<4xf32>) { + %values = func.call @ggml_q8_0_f32_vector4(%weight, %row_byte_base, %q8_block, %packet) : (buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %q8_1_values = scf.if %is_q8_1 -> (vector<4xf32>) { + %values = func.call @ggml_q8_1_f32_vector4(%weight, %row_byte_base, %q8_block, %packet) : (buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %f16_values = scf.if %is_f16 -> (vector<4xf32>) { + %values = func.call @ggml_f16_f32_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %bf16_values = scf.if %is_bf16 -> (vector<4xf32>) { + %values = func.call @ggml_bf16_f32_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %f32_values = scf.if %is_f32 -> (vector<4xf32>) { + %values = func.call @ggml_f32_f32_vector4(%weight, %row_byte_base, %input_size, %k) : (buffer, offset, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %selected_q1_0 = scf.select %is_q1_0, %q1_0_values, %zero : vector<4xf32> + %selected_q3 = scf.select %is_q3, %q3_values, %selected_q1_0 : vector<4xf32> + %selected_q4 = scf.select %is_q4, %q4_values, %selected_q3 : vector<4xf32> + %selected_q5 = scf.select %is_q5, %q5_values, %selected_q4 : vector<4xf32> + %selected_q6 = scf.select %is_q6, %q6_values, %selected_q5 : vector<4xf32> + %selected_q4_0 = scf.select %is_q4_0, %q4_0_values, %selected_q6 : vector<4xf32> + %selected_q4_1 = scf.select %is_q4_1, %q4_1_values, %selected_q4_0 : vector<4xf32> + %selected_q5_0 = scf.select %is_q5_0, %q5_0_values, %selected_q4_1 : vector<4xf32> + %selected_q5_1 = scf.select %is_q5_1, %q5_1_values, %selected_q5_0 : vector<4xf32> + %selected_iq2_s = scf.select %is_iq2_s, %iq2_s_values, %selected_q5_1 : vector<4xf32> + %selected_iq4_nl = scf.select %is_iq4_nl, %iq4_nl_values, %selected_iq2_s : vector<4xf32> + %selected_iq3_s = scf.select %is_iq3_s, %iq3_s_values, %selected_iq4_nl : vector<4xf32> + %selected_iq4_xs = scf.select %is_iq4_xs, %iq4_xs_values, %selected_iq3_s : vector<4xf32> + %selected_q8_0 = scf.select %is_q8_0, %q8_0_values, %selected_iq4_xs : vector<4xf32> + %selected_q8_1 = scf.select %is_q8_1, %q8_1_values, %selected_q8_0 : vector<4xf32> + %selected_f16 = scf.select %is_f16, %f16_values, %selected_q8_1 : vector<4xf32> + %selected_bf16 = scf.select %is_bf16, %bf16_values, %selected_f16 : vector<4xf32> + %selected = scf.select %is_f32, %f32_values, %selected_bf16 : vector<4xf32> + func.return %selected : vector<4xf32> +} + +// Reuse immutable Q5 metadata across all four 64-element quarters. +func.def inline @ggml_q5k_load_decode_state(%weight: buffer, %row_base: offset, %block: index, %packet: index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 176 : offset + %qh_offset = index.constant 16 : offset + %block_add = index.scale %block, %block_bytes : index, offset -> offset + %block_base = index.add %row_base, %block_add : offset + %qh_base = index.add %block_base, %qh_offset : offset + %header_view = buffer.view %weight[%block_base] : buffer -> view<4xi32> + %header = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> + %dm_word = vector.extract %header[0] : vector<4xi32> -> i32 + %dm_vector = vector.from_elements %dm_word : vector<1xi32> + %dm = vector.bitcast %dm_vector : vector<1xi32> to vector<2xf16> + %s0 = vector.extract %header[1] : vector<4xi32> -> i32 + %s1 = vector.extract %header[2] : vector<4xi32> -> i32 + %s2 = vector.extract %header[3] : vector<4xi32> -> i32 + %scales = vector.from_elements %s0, %s1, %s2 : vector<3xi32> + %qh_view = buffer.view %weight[%qh_base] : buffer -> view<32xi8> + %packet_base = index.mul %packet, %c4 : index + %qh = vector.load %qh_view[%packet_base] : view<32xi8> -> vector<4xi8> + func.return %dm, %scales, %qh : vector<2xf16>, vector<3xi32>, vector<4xi8> +} + +func.def inline @ggml_q5k_preloaded_f32_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>, %qh_bytes: vector<4xi8>) -> (vector<4xf32>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 176 : offset + %code_offset = index.constant 48 : offset + %c1_i32 = scalar.constant 1 : i32 + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q5_group [range(%q5_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q5_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %qh_shift_i32 = index.cast %bounded_group : index to i32 + %qh_bit = scalar.shli %c1_i32, %qh_shift_i32 : i32 + %qh0 = vector.extract %qh_bytes[0] : vector<4xi8> -> i8 + %qh1 = vector.extract %qh_bytes[1] : vector<4xi8> -> i8 + %qh2 = vector.extract %qh_bytes[2] : vector<4xi8> -> i8 + %qh3 = vector.extract %qh_bytes[3] : vector<4xi8> -> i8 + %high0 = func.call @ggml_q5k_high_bit_f32(%qh0, %qh_bit) : (i8, i32) -> (f32) + %high1 = func.call @ggml_q5k_high_bit_f32(%qh1, %qh_bit) : (i8, i32) -> (f32) + %high2 = func.call @ggml_q5k_high_bit_f32(%qh2, %qh_bit) : (i8, i32) -> (f32) + %high3 = func.call @ggml_q5k_high_bit_f32(%qh3, %qh_bit) : (i8, i32) -> (f32) + %q0_low = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1_low = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2_low = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3_low = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %q0 = scalar.addf %q0_low, %high0 : f32 + %q1 = scalar.addf %q1_low, %high1 : f32 + %q2 = scalar.addf %q2_low, %high2 : f32 + %q3 = scalar.addf %q3_low, %high3 : f32 + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + +// Reuse the IQ4 header across all four 64-element quarters. +func.def inline @ggml_iq4xs_load_decode_state(%weight: buffer, %row_base: offset, %block: index) -> (i32, i32) { + %c0 = index.constant 0 : index + %block_bytes = index.constant 136 : offset + %block_add = index.scale %block, %block_bytes : index, offset -> offset + %block_base = index.add %row_base, %block_add : offset + %header_view = buffer.view %weight[%block_base] : buffer -> view<2xi32> + %header = vector.load %header_view[%c0] : view<2xi32> -> vector<2xi32> + %high = vector.extract %header[0] : vector<2xi32> -> i32 + %low = vector.extract %header[1] : vector<2xi32> -> i32 + func.return %high, %low : i32, i32 +} + +func.def inline @ggml_iq4xs_preloaded_f32_vector4(%iq4nl_table: vector<16xi8>, %weight: buffer, %row_byte_base: offset, %iq4_block: index, %iq4_group: index, %packet: index, %header: i32, %scales_l: i32) -> (vector<4xf32>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 136 : offset + %code_offset = index.constant 8 : offset + %c3_i32 = scalar.constant 3 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c16_i32 = scalar.constant 16 : i32 + %c32_i32 = scalar.constant 32 : i32 + %c65535_i32 = scalar.constant 65535 : i32 + %bounded_group = index.assume %iq4_group [range(%iq4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %iq4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_vector = vector.from_elements %header : vector<1xi32> + %header_halves = vector.bitcast %header_vector : vector<1xi32> to vector<2xf16> + %d_f16 = vector.extract %header_halves[0] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %scales_h_shifted = scalar.shrui %header, %c16_i32 : i32 + %scales_h = scalar.andi %scales_h_shifted, %c65535_i32 : i32 + %scale_l_shift_index = index.mul %bounded_group, %c4 : index + %scale_l_shift = index.cast %scale_l_shift_index : index to i32 + %scale_l_shifted = scalar.shrui %scales_l, %scale_l_shift : i32 + %scale_l = scalar.andi %scale_l_shifted, %c15_i32 : i32 + %scale_h_shift_index = index.mul %bounded_group, %c2 : index + %scale_h_shift = index.cast %scale_h_shift_index : index to i32 + %scale_h_shifted = scalar.shrui %scales_h, %scale_h_shift : i32 + %scale_h = scalar.andi %scale_h_shifted, %c3_i32 : i32 + %scale_h_high = scalar.shli %scale_h, %c4_i32 : i32 + %local_scale = scalar.ori %scale_l, %scale_h_high : i32 + %centered_scale = scalar.subi %local_scale, %c32_i32 : i32 + %centered_scale_f32 = scalar.sitofp %centered_scale : i32 to f32 + %block_scale = scalar.mulf %d, %centered_scale_f32 : f32 + %packet_page = index.rem %bounded_packet, %c4 : index + %q_word_base = index.mul %bounded_group, %c4 : index + %q_word_index0 = index.add %q_word_base, %packet_page : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %uses_high = index.cmp uge, %bounded_packet, %c4 : index + %codes = func.call @ggml_iq4xs_table_codes4(%q_word, %uses_high) : (vector<1xi32>, i1) -> (vector<4xi8>) + %lookup_i8 = func.call @ggml_iq4nl_lookup_i8_vector4(%iq4nl_table, %codes) : (vector<16xi8>, vector<4xi8>) -> (vector<4xi8>) + %lookup0_i8 = vector.extract %lookup_i8[0] : vector<4xi8> -> i8 + %lookup1_i8 = vector.extract %lookup_i8[1] : vector<4xi8> -> i8 + %lookup2_i8 = vector.extract %lookup_i8[2] : vector<4xi8> -> i8 + %lookup3_i8 = vector.extract %lookup_i8[3] : vector<4xi8> -> i8 + %lookup0_i32 = scalar.extsi %lookup0_i8 : i8 to i32 + %lookup1_i32 = scalar.extsi %lookup1_i8 : i8 to i32 + %lookup2_i32 = scalar.extsi %lookup2_i8 : i8 to i32 + %lookup3_i32 = scalar.extsi %lookup3_i8 : i8 to i32 + %lookup0 = scalar.sitofp %lookup0_i32 : i32 to f32 + %lookup1 = scalar.sitofp %lookup1_i32 : i32 to f32 + %lookup2 = scalar.sitofp %lookup2_i32 : i32 to f32 + %lookup3 = scalar.sitofp %lookup3_i32 : i32 to f32 + %value0 = scalar.mulf %block_scale, %lookup0 : f32 + %value1 = scalar.mulf %block_scale, %lookup1 : f32 + %value2 = scalar.mulf %block_scale, %lookup2 : f32 + %value3 = scalar.mulf %block_scale, %lookup3 : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/binary_apply.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/binary_apply.loom new file mode 100644 index 000000000000..bdbb77a86432 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/binary_apply.loom @@ -0,0 +1,141 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +template.def<@ggml.binary_f32.apply> device priority(1) @ggml_binary_f32_apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) { + %c0_f32 = scalar.constant 0.0 : f32 + %c0_5_f32 = scalar.constant 0.5 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c2_f32 = scalar.constant 2.0 : f32 + %cn2_f32 = scalar.constant -2.0 : f32 + %gelu_coef = scalar.constant 0.044715 : f32 + %gelu_quick_coef = scalar.constant -1.702 : f32 + %sqrt_2_over_pi = scalar.constant 0.79788456080286544 : f32 + %sqrt_2_inv = scalar.constant 0.70710678118654745 : f32 + + // ADD (op 0): x + y. + %op_add = index.constant 0 : index + %sum = scalar.addf %lhs, %rhs : f32 + %is_add = index.cmp eq, %op, %op_add : index + %add_selected = scf.select %is_add, %sum, %lhs : f32 + + // SUB (op 1): x - y. + %op_sub = index.constant 1 : index + %difference = scalar.subf %lhs, %rhs : f32 + %is_sub = index.cmp eq, %op, %op_sub : index + %sub_selected = scf.select %is_sub, %difference, %add_selected : f32 + + // MUL (op 2): x * y. + %op_mul = index.constant 2 : index + %product = scalar.mulf %lhs, %rhs : f32 + %is_mul = index.cmp eq, %op, %op_mul : index + %mul_selected = scf.select %is_mul, %product, %sub_selected : f32 + + // DIV (op 3): x / y. + %op_div = index.constant 3 : index + %quotient = scalar.divf %lhs, %rhs : f32 + %is_div = index.cmp eq, %op, %op_div : index + %div_selected = scf.select %is_div, %quotient, %mul_selected : f32 + + // SWIGLU (op 4): silu(x) * y. + %op_swiglu = index.constant 4 : index + %silu_negative = scalar.negf %lhs : f32 + %silu_exp = scalar.expf %silu_negative : f32 + %silu_denominator = scalar.addf %c1_f32, %silu_exp : f32 + %silu_sigmoid = scalar.divf %c1_f32, %silu_denominator : f32 + %silu = scalar.mulf %lhs, %silu_sigmoid : f32 + %swiglu = scalar.mulf %silu, %rhs : f32 + %is_swiglu = index.cmp eq, %op, %op_swiglu : index + %swiglu_selected = scf.select %is_swiglu, %swiglu, %div_selected : f32 + + // GEGLU (op 5): gelu(x) * y. + %op_geglu = index.constant 5 : index + %gelu_x2 = scalar.mulf %lhs, %lhs : f32 + %gelu_poly0 = scalar.mulf %gelu_coef, %gelu_x2 : f32 + %gelu_poly = scalar.addf %c1_f32, %gelu_poly0 : f32 + %gelu_inner0 = scalar.mulf %lhs, %gelu_poly : f32 + %gelu_inner = scalar.mulf %sqrt_2_over_pi, %gelu_inner0 : f32 + %gelu_tanh_scaled = scalar.mulf %cn2_f32, %gelu_inner : f32 + %gelu_tanh_exp = scalar.expf %gelu_tanh_scaled : f32 + %gelu_tanh_denominator = scalar.addf %c1_f32, %gelu_tanh_exp : f32 + %gelu_tanh_ratio = scalar.divf %c2_f32, %gelu_tanh_denominator : f32 + %gelu_tanh = scalar.subf %gelu_tanh_ratio, %c1_f32 : f32 + %gelu_one_plus = scalar.addf %c1_f32, %gelu_tanh : f32 + %gelu_half_x = scalar.mulf %c0_5_f32, %lhs : f32 + %gelu = scalar.mulf %gelu_half_x, %gelu_one_plus : f32 + %geglu = scalar.mulf %gelu, %rhs : f32 + %is_geglu = index.cmp eq, %op, %op_geglu : index + %geglu_selected = scf.select %is_geglu, %geglu, %swiglu_selected : f32 + + // REGLU (op 6): relu(x) * y. + %op_reglu = index.constant 6 : index + %relu_is_positive = scalar.cmpf ogt, %lhs, %c0_f32 : f32 + %relu = scf.select %relu_is_positive, %lhs, %c0_f32 : f32 + %reglu = scalar.mulf %relu, %rhs : f32 + %is_reglu = index.cmp eq, %op, %op_reglu : index + %reglu_selected = scf.select %is_reglu, %reglu, %geglu_selected : f32 + + // GEGLU_ERF (op 7): gelu_erf(x) * y. + %op_geglu_erf = index.constant 7 : index + %gelu_erf_scaled = scalar.mulf %sqrt_2_inv, %lhs : f32 + %gelu_erf_value = scalar.erff %gelu_erf_scaled : f32 + %gelu_erf_one_plus = scalar.addf %c1_f32, %gelu_erf_value : f32 + %gelu_erf_half_x = scalar.mulf %c0_5_f32, %lhs : f32 + %gelu_erf = scalar.mulf %gelu_erf_half_x, %gelu_erf_one_plus : f32 + %geglu_erf = scalar.mulf %gelu_erf, %rhs : f32 + %is_geglu_erf = index.cmp eq, %op, %op_geglu_erf : index + %geglu_erf_selected = scf.select %is_geglu_erf, %geglu_erf, %reglu_selected : f32 + + // GEGLU_QUICK (op 8): gelu_quick(x) * y. + %op_geglu_quick = index.constant 8 : index + %gelu_quick_scaled = scalar.mulf %gelu_quick_coef, %lhs : f32 + %gelu_quick_exp = scalar.expf %gelu_quick_scaled : f32 + %gelu_quick_denominator = scalar.addf %c1_f32, %gelu_quick_exp : f32 + %gelu_quick_sigmoid = scalar.divf %c1_f32, %gelu_quick_denominator : f32 + %gelu_quick = scalar.mulf %lhs, %gelu_quick_sigmoid : f32 + %geglu_quick = scalar.mulf %gelu_quick, %rhs : f32 + %is_geglu_quick = index.cmp eq, %op, %op_geglu_quick : index + %result = scf.select %is_geglu_quick, %geglu_quick, %geglu_erf_selected : f32 + template.return %result : f32 +} + +template.decl @ggml.binary_f32.apply_vector8(%op: index, %lhs: vector<8xf32>, %rhs: vector<8xf32>) -> (vector<8xf32>) + +template.def<@ggml.binary_f32.apply_vector8> device @ggml_binary_f32_apply_vector8(%op: index, %lhs: vector<8xf32>, %rhs: vector<8xf32>) -> (vector<8xf32>) { + %op_swiglu = index.constant 4 : index + %is_swiglu = index.cmp eq, %op, %op_swiglu : index + %result = scf.if %is_swiglu -> (vector<8xf32>) { + %silu = vector.siluf %lhs : vector<8xf32> + %value = vector.mulf %silu, %rhs : vector<8xf32> + scf.yield %value : vector<8xf32> + } else { + %lhs0 = vector.extract %lhs[0] : vector<8xf32> -> f32 + %rhs0 = vector.extract %rhs[0] : vector<8xf32> -> f32 + %value0 = template.apply<@ggml.binary_f32.apply>(%op, %lhs0, %rhs0) : (index, f32, f32) -> (f32) + %lhs1 = vector.extract %lhs[1] : vector<8xf32> -> f32 + %rhs1 = vector.extract %rhs[1] : vector<8xf32> -> f32 + %value1 = template.apply<@ggml.binary_f32.apply>(%op, %lhs1, %rhs1) : (index, f32, f32) -> (f32) + %lhs2 = vector.extract %lhs[2] : vector<8xf32> -> f32 + %rhs2 = vector.extract %rhs[2] : vector<8xf32> -> f32 + %value2 = template.apply<@ggml.binary_f32.apply>(%op, %lhs2, %rhs2) : (index, f32, f32) -> (f32) + %lhs3 = vector.extract %lhs[3] : vector<8xf32> -> f32 + %rhs3 = vector.extract %rhs[3] : vector<8xf32> -> f32 + %value3 = template.apply<@ggml.binary_f32.apply>(%op, %lhs3, %rhs3) : (index, f32, f32) -> (f32) + %lhs4 = vector.extract %lhs[4] : vector<8xf32> -> f32 + %rhs4 = vector.extract %rhs[4] : vector<8xf32> -> f32 + %value4 = template.apply<@ggml.binary_f32.apply>(%op, %lhs4, %rhs4) : (index, f32, f32) -> (f32) + %lhs5 = vector.extract %lhs[5] : vector<8xf32> -> f32 + %rhs5 = vector.extract %rhs[5] : vector<8xf32> -> f32 + %value5 = template.apply<@ggml.binary_f32.apply>(%op, %lhs5, %rhs5) : (index, f32, f32) -> (f32) + %lhs6 = vector.extract %lhs[6] : vector<8xf32> -> f32 + %rhs6 = vector.extract %rhs[6] : vector<8xf32> -> f32 + %value6 = template.apply<@ggml.binary_f32.apply>(%op, %lhs6, %rhs6) : (index, f32, f32) -> (f32) + %lhs7 = vector.extract %lhs[7] : vector<8xf32> -> f32 + %rhs7 = vector.extract %rhs[7] : vector<8xf32> -> f32 + %value7 = template.apply<@ggml.binary_f32.apply>(%op, %lhs7, %rhs7) : (index, f32, f32) -> (f32) + %value = vector.from_elements %value0, %value1, %value2, %value3, %value4, %value5, %value6, %value7 : vector<8xf32> + scf.yield %value : vector<8xf32> + } + template.return %result : vector<8xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/broadcast_index.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/broadcast_index.loom new file mode 100644 index 000000000000..46096a7c04fe --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/broadcast_index.loom @@ -0,0 +1,33 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.elementwise.broadcast_index(%linear: index, %n0: index, %n1: index, %n2: index, %coord0: index, %coord1: index, %coord2: index, %coord3: index, %broadcast_dim0_flag: index, %broadcast_dim1_flag: index, %broadcast_dim2_flag: index, %broadcast_dim3_flag: index, %zero: index, %one: index) -> (index) + +template.def<@ggml.elementwise.broadcast_index> device @ggml_elementwise_broadcast_index(%linear: index, %n0: index, %n1: index, %n2: index, %coord0: index, %coord1: index, %coord2: index, %coord3: index, %broadcast_dim0_flag: index, %broadcast_dim1_flag: index, %broadcast_dim2_flag: index, %broadcast_dim3_flag: index, %zero: index, %one: index) -> (index) { + %broadcast_dims01 = index.add %broadcast_dim0_flag, %broadcast_dim1_flag : index + %broadcast_dims23 = index.add %broadcast_dim2_flag, %broadcast_dim3_flag : index + %broadcast_dims = index.add %broadcast_dims01, %broadcast_dims23 : index + %no_broadcast = index.cmp eq, %broadcast_dims, %zero : index + %broadcast_dim0 = index.cmp eq, %broadcast_dim0_flag, %one : index + %broadcast_dim1 = index.cmp eq, %broadcast_dim1_flag, %one : index + %broadcast_dim2 = index.cmp eq, %broadcast_dim2_flag, %one : index + %broadcast_dim3 = index.cmp eq, %broadcast_dim3_flag, %one : index + %source_ne0 = scf.select %broadcast_dim0, %one, %n0 : index + %source_ne1 = scf.select %broadcast_dim1, %one, %n1 : index + %source_ne2 = scf.select %broadcast_dim2, %one, %n2 : index + %source_coord0 = scf.select %broadcast_dim0, %zero, %coord0 : index + %source_coord1 = scf.select %broadcast_dim1, %zero, %coord1 : index + %source_coord2 = scf.select %broadcast_dim2, %zero, %coord2 : index + %source_coord3 = scf.select %broadcast_dim3, %zero, %coord3 : index + %source_stride1 = index.add %source_ne0, %zero : index + %source_stride2 = index.mul %source_stride1, %source_ne1 : index + %source_stride3 = index.mul %source_stride2, %source_ne2 : index + %source_term1 = index.mul %source_coord1, %source_stride1 : index + %source_term2 = index.mul %source_coord2, %source_stride2 : index + %source_term3 = index.mul %source_coord3, %source_stride3 : index + %source_linear01 = index.add %source_coord0, %source_term1 : index + %source_linear012 = index.add %source_linear01, %source_term2 : index + %broadcast_linear = index.add %source_linear012, %source_term3 : index + %source_linear = scf.select %no_broadcast, %linear, %broadcast_linear : index + template.return %source_linear : index +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/scale_apply.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/scale_apply.loom new file mode 100644 index 000000000000..ea648f3a0b5c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/scale_apply.loom @@ -0,0 +1,10 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.elementwise.scale_apply(%value: f32, %scale: f32, %bias: f32) -> (f32) + +template.def<@ggml.elementwise.scale_apply> device @ggml_elementwise_scale_apply(%value: f32, %scale: f32, %bias: f32) -> (f32) { + %scaled = scalar.mulf %value, %scale : f32 + %result = scalar.addf %scaled, %bias : f32 + template.return %result : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/symmetric_i4_binary.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/symmetric_i4_binary.loom new file mode 100644 index 000000000000..248537d8a40a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/symmetric_i4_binary.loom @@ -0,0 +1,133 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +config.decl @ggml.binary_swiglu_symmetric_i4.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 64)] + +config.decl @ggml.binary_swiglu_symmetric_i4.token_count : %value: index where [range(%value, 1, 16)] + +template.decl @ggml.binary_swiglu_symmetric_i4_k32.body(%lhs: buffer, %rhs: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.def<@ggml.binary_swiglu_symmetric_i4_k32.body> device @ggml_binary_swiglu_symmetric_i4_k32_body(%lhs: buffer, %rhs: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %amax_epsilon = scalar.constant 1.0000000000000001e-30 : f32 + %one_seventh = scalar.constant 0.14285714285714285 : f32 + %seven = scalar.constant 7.0 : f32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %shuffle_width = scalar.constant 32 : i32 + %shift4 = scalar.constant 4 : i32 + %shift8 = scalar.constant 8 : i32 + %mask15 = vector.constant 15 : vector<2xi32> + + %input_size0 = config.get @ggml.binary_swiglu_symmetric_i4.input_size : index + %token_count0 = config.get @ggml.binary_swiglu_symmetric_i4.token_count : index + %input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 64)] : index + %token_count = index.assume %token_count0 [range(%token_count0, 1, 16)] : index + %element_count0 = index.mul %input_size, %token_count : index + %element_count = index.assume %element_count0 [range(%element_count0, 256, 524288), mul(%element_count0, 64)] : index + %groups32 = index.div %element_count, %c32 : index + %groups64 = index.div %element_count, %c64 : index + %qs_halfwords = index.div %element_count, %c4 : index + + %workgroup = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index + %group_base = index.mul %workgroup, %c8 : index + %group64 = index.add %group_base, %subgroup : index + %valid_group = index.cmp ult, %group64, %groups64 : index + + %lhs_global = buffer.assume.memory_space %lhs : buffer + %rhs_global = buffer.assume.memory_space %rhs : buffer + %output_global = buffer.assume.memory_space %output : buffer + %i4_qs_global = buffer.assume.memory_space %i4_qs : buffer + %i4_ds_global = buffer.assume.memory_space %i4_ds : buffer + %i4_sums_global = buffer.assume.memory_space %i4_sums : buffer + %lhs_na, %rhs_na, %output_na, %i4_qs_na, %i4_ds_na, %i4_sums_na = buffer.assume.noalias %lhs_global, %rhs_global, %output_global, %i4_qs_global, %i4_ds_global, %i4_sums_global : buffer, buffer, buffer, buffer, buffer, buffer + %lhs_view = buffer.view %lhs_na[%base] : buffer -> view<[%element_count]xf32> + %rhs_view = buffer.view %rhs_na[%base] : buffer -> view<[%element_count]xf32> + %output_view = buffer.view %output_na[%base] : buffer -> view<[%element_count]xf32> + %i4_qs_view = buffer.view %i4_qs_na[%base] : buffer -> view<[%qs_halfwords]xi16> + %i4_ds_view = buffer.view %i4_ds_na[%base] : buffer -> view<[%groups32]xf32> + %i4_sums_view = buffer.view %i4_sums_na[%base] : buffer -> view<[%groups32]xi32> + + scf.if %valid_group { + %group_element_base = index.mul %group64, %c64 : index + %lane_element_offset = index.mul %lane, %c2 : index + %linear0 = index.add %group_element_base, %lane_element_offset : index + %linear = index.assume %linear0 [range(%linear0, 0, 524286), mul(%linear0, 2)] : index + %lhs_values = vector.load %lhs_view[%linear] : view<[%element_count]xf32> -> vector<2xf32> + %rhs_values = vector.load %rhs_view[%linear] : view<[%element_count]xf32> -> vector<2xf32> + %activated = vector.siluf %lhs_values : vector<2xf32> + %result = vector.mulf %activated, %rhs_values : vector<2xf32> + vector.store %result, %output_view[%linear] : vector<2xf32>, view<[%element_count]xf32> + + %absolute_values = vector.absf %result : vector<2xf32> + %lane_max = vector.reduce %absolute_values, %c0_f32 : vector<2xf32>, f32 + %group_max = kernel.subgroup.reduce %lane_max : f32 + %amax = scalar.maxnumf %group_max, %amax_epsilon : f32 + %scale = scalar.mulf %amax, %one_seventh : f32 + %is_scale_leader = index.cmp eq, %lane, %c0 : index + %group32_base = index.mul %group64, %c2 : index + %group32_high = index.add %group32_base, %c1 : index + scf.if %is_scale_leader { + view.store %scale, %i4_ds_view[%group32_base] : f32, view<[%groups32]xf32> + view.store %scale, %i4_ds_view[%group32_high] : f32, view<[%groups32]xf32> + } + + %rscale = scalar.divf %seven, %amax : f32 + %rscale_vector = vector.splat %rscale : vector<2xf32> + %scaled = vector.mulf %result, %rscale_vector : vector<2xf32> + %rounded = vector.roundf %scaled : vector<2xf32> + %quantized = vector.fptosi %rounded : vector<2xf32> to vector<2xi32> + %nibbles = vector.andi %quantized, %mask15 : vector<2xi32> + %q0 = vector.extract %nibbles[0] : vector<2xi32> -> i32 + %q1 = vector.extract %nibbles[1] : vector<2xi32> -> i32 + %q1_shifted = scalar.shli %q1, %shift4 : i32 + %packed_pair = scalar.ori %q0, %q1_shifted : i32 + %peer_pair, %peer_valid = kernel.subgroup.shuffle %packed_pair, %xor1, %shuffle_width : i32, i32, i32 + %peer_shifted = scalar.shli %peer_pair, %shift8 : i32 + %packed_word = scalar.ori %packed_pair, %peer_shifted : i32 + %packed_i16 = scalar.trunci %packed_word : i32 to i16 + %lane_pair = index.div %lane, %c2 : index + %word_base = index.mul %group64, %c16 : index + %word_index = index.add %word_base, %lane_pair : index + %lane_in_pair = index.rem %lane, %c2 : index + %is_word_leader = index.cmp eq, %lane_in_pair, %c0 : index + scf.if %is_word_leader { + view.store %packed_i16, %i4_qs_view[%word_index] : i16, view<[%qs_halfwords]xi16> + } + + %lane_sum = vector.reduce %rounded, %c0_f32 : vector<2xf32>, f32 + %sum_x1_peer, %sum_x1_valid = kernel.subgroup.shuffle %lane_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum_x1 = scalar.addf %lane_sum, %sum_x1_peer : f32 + %sum_x2_peer, %sum_x2_valid = kernel.subgroup.shuffle %sum_x1, %xor2, %shuffle_width : f32, i32, i32 + %sum_x2 = scalar.addf %sum_x1, %sum_x2_peer : f32 + %sum_x4_peer, %sum_x4_valid = kernel.subgroup.shuffle %sum_x2, %xor4, %shuffle_width : f32, i32, i32 + %sum_x4 = scalar.addf %sum_x2, %sum_x4_peer : f32 + %sum_x8_peer, %sum_x8_valid = kernel.subgroup.shuffle %sum_x4, %xor8, %shuffle_width : f32, i32, i32 + %group_sum = scalar.addf %sum_x4, %sum_x8_peer : f32 + %group_half = index.div %lane, %c16 : index + %lane_in_half = index.rem %lane, %c16 : index + %is_sum_leader = index.cmp eq, %lane_in_half, %c0 : index + scf.if %is_sum_leader { + %sum_index = index.add %group32_base, %group_half : index + %sum_value = scalar.fptosi %group_sum : f32 to i32 + view.store %sum_value, %i4_sums_view[%sum_index] : i32, view<[%groups32]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/unary_apply.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/unary_apply.loom new file mode 100644 index 000000000000..84b1d30a768d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/elementwise/unary_apply.loom @@ -0,0 +1,380 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.unary_f32.apply(%op: index, %value: f32) -> (f32) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +template.def<@ggml.unary_f32.apply> device priority(1) @ggml_unary_f32_apply(%op: index, %value: f32) -> (f32) { + %c0_f32 = scalar.constant 0.0 : f32 + %c0_5_f32 = scalar.constant 0.5 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c2_f32 = scalar.constant 2.0 : f32 + %c3_f32 = scalar.constant 3.0 : f32 + %c6_f32 = scalar.constant 6.0 : f32 + %cn1_f32 = scalar.constant -1.0 : f32 + %cn2_f32 = scalar.constant -2.0 : f32 + %gelu_coef = scalar.constant 0.044715 : f32 + %gelu_quick_coef = scalar.constant -1.702 : f32 + %sqrt_2_over_pi = scalar.constant 0.79788456080286544 : f32 + %sqrt_2_inv = scalar.constant 0.70710678118654745 : f32 + + %op_neg = index.constant 0 : index + %negated = scalar.negf %value : f32 + %is_neg = index.cmp eq, %op, %op_neg : index + %neg_selected = scf.select %is_neg, %negated, %value : f32 + + %op_abs = index.constant 1 : index + %abs_is_negative = scalar.cmpf olt, %value, %c0_f32 : f32 + %absolute = scf.select %abs_is_negative, %negated, %value : f32 + %is_abs = index.cmp eq, %op, %op_abs : index + %abs_selected = scf.select %is_abs, %absolute, %neg_selected : f32 + + %op_relu = index.constant 2 : index + %relu_is_positive = scalar.cmpf ogt, %value, %c0_f32 : f32 + %relu = scf.select %relu_is_positive, %value, %c0_f32 : f32 + %is_relu = index.cmp eq, %op, %op_relu : index + %relu_selected = scf.select %is_relu, %relu, %abs_selected : f32 + + %op_step = index.constant 3 : index + %step_is_positive = scalar.cmpf ogt, %value, %c0_f32 : f32 + %step = scf.select %step_is_positive, %c1_f32, %c0_f32 : f32 + %is_step = index.cmp eq, %op, %op_step : index + %step_selected = scf.select %is_step, %step, %relu_selected : f32 + + %op_sqr = index.constant 4 : index + %squared = scalar.mulf %value, %value : f32 + %is_sqr = index.cmp eq, %op, %op_sqr : index + %sqr_selected = scf.select %is_sqr, %squared, %step_selected : f32 + + %op_sgn = index.constant 5 : index + %sgn_is_negative = scalar.cmpf olt, %value, %c0_f32 : f32 + %sgn_is_positive = scalar.cmpf ogt, %value, %c0_f32 : f32 + %negative_sgn = scf.select %sgn_is_negative, %cn1_f32, %c0_f32 : f32 + %sgn = scf.select %sgn_is_positive, %c1_f32, %negative_sgn : f32 + %is_sgn = index.cmp eq, %op, %op_sgn : index + %sgn_selected = scf.select %is_sgn, %sgn, %sqr_selected : f32 + + %op_floor = index.constant 6 : index + %floored = scalar.floorf %value : f32 + %is_floor = index.cmp eq, %op, %op_floor : index + %floor_selected = scf.select %is_floor, %floored, %sgn_selected : f32 + + %op_ceil = index.constant 7 : index + %ceiled = scalar.ceilf %value : f32 + %is_ceil = index.cmp eq, %op, %op_ceil : index + %ceil_selected = scf.select %is_ceil, %ceiled, %floor_selected : f32 + + %op_round = index.constant 8 : index + %rounded = scalar.roundf %value : f32 + %is_round = index.cmp eq, %op, %op_round : index + %round_selected = scf.select %is_round, %rounded, %ceil_selected : f32 + + %op_trunc = index.constant 9 : index + %truncated = scalar.truncf %value : f32 + %is_trunc = index.cmp eq, %op, %op_trunc : index + %trunc_selected = scf.select %is_trunc, %truncated, %round_selected : f32 + + %op_tanh = index.constant 10 : index + %tanh_scaled = scalar.mulf %cn2_f32, %value : f32 + %tanh_exp = scalar.expf %tanh_scaled : f32 + %tanh_denominator = scalar.addf %c1_f32, %tanh_exp : f32 + %tanh_ratio = scalar.divf %c2_f32, %tanh_denominator : f32 + %tanh = scalar.subf %tanh_ratio, %c1_f32 : f32 + %is_tanh = index.cmp eq, %op, %op_tanh : index + %tanh_selected = scf.select %is_tanh, %tanh, %trunc_selected : f32 + + %op_elu = index.constant 11 : index + %elu_is_positive = scalar.cmpf ogt, %value, %c0_f32 : f32 + %elu_exp = scalar.expf %value : f32 + %elu_expm1 = scalar.subf %elu_exp, %c1_f32 : f32 + %elu = scf.select %elu_is_positive, %value, %elu_expm1 : f32 + %is_elu = index.cmp eq, %op, %op_elu : index + %elu_selected = scf.select %is_elu, %elu, %tanh_selected : f32 + + %op_sigmoid = index.constant 12 : index + %sigmoid_negative = scalar.negf %value : f32 + %sigmoid_exp = scalar.expf %sigmoid_negative : f32 + %sigmoid_denominator = scalar.addf %c1_f32, %sigmoid_exp : f32 + %sigmoid = scalar.divf %c1_f32, %sigmoid_denominator : f32 + %is_sigmoid = index.cmp eq, %op, %op_sigmoid : index + %sigmoid_selected = scf.select %is_sigmoid, %sigmoid, %elu_selected : f32 + + %op_gelu = index.constant 13 : index + %gelu_x2 = scalar.mulf %value, %value : f32 + %gelu_poly0 = scalar.mulf %gelu_coef, %gelu_x2 : f32 + %gelu_poly = scalar.addf %c1_f32, %gelu_poly0 : f32 + %gelu_inner0 = scalar.mulf %value, %gelu_poly : f32 + %gelu_inner = scalar.mulf %sqrt_2_over_pi, %gelu_inner0 : f32 + %gelu_tanh_scaled = scalar.mulf %cn2_f32, %gelu_inner : f32 + %gelu_tanh_exp = scalar.expf %gelu_tanh_scaled : f32 + %gelu_tanh_denominator = scalar.addf %c1_f32, %gelu_tanh_exp : f32 + %gelu_tanh_ratio = scalar.divf %c2_f32, %gelu_tanh_denominator : f32 + %gelu_tanh = scalar.subf %gelu_tanh_ratio, %c1_f32 : f32 + %gelu_one_plus = scalar.addf %c1_f32, %gelu_tanh : f32 + %gelu_half_x = scalar.mulf %c0_5_f32, %value : f32 + %gelu = scalar.mulf %gelu_half_x, %gelu_one_plus : f32 + %is_gelu = index.cmp eq, %op, %op_gelu : index + %gelu_selected = scf.select %is_gelu, %gelu, %sigmoid_selected : f32 + + %op_gelu_quick = index.constant 14 : index + %gelu_quick_scaled = scalar.mulf %gelu_quick_coef, %value : f32 + %gelu_quick_exp = scalar.expf %gelu_quick_scaled : f32 + %gelu_quick_denominator = scalar.addf %c1_f32, %gelu_quick_exp : f32 + %gelu_quick_sigmoid = scalar.divf %c1_f32, %gelu_quick_denominator : f32 + %gelu_quick = scalar.mulf %value, %gelu_quick_sigmoid : f32 + %is_gelu_quick = index.cmp eq, %op, %op_gelu_quick : index + %gelu_quick_selected = scf.select %is_gelu_quick, %gelu_quick, %gelu_selected : f32 + + %op_silu = index.constant 15 : index + %silu_negative = scalar.negf %value : f32 + %silu_exp = scalar.expf %silu_negative : f32 + %silu_denominator = scalar.addf %c1_f32, %silu_exp : f32 + %silu_sigmoid = scalar.divf %c1_f32, %silu_denominator : f32 + %silu = scalar.mulf %value, %silu_sigmoid : f32 + %is_silu = index.cmp eq, %op, %op_silu : index + %silu_selected = scf.select %is_silu, %silu, %gelu_quick_selected : f32 + + %op_hardswish = index.constant 16 : index + %hardswish_bias = scalar.addf %value, %c3_f32 : f32 + %hardswish_scaled = scalar.divf %hardswish_bias, %c6_f32 : f32 + %hardswish_above_zero = scalar.cmpf ogt, %hardswish_scaled, %c0_f32 : f32 + %hardswish_low = scf.select %hardswish_above_zero, %hardswish_scaled, %c0_f32 : f32 + %hardswish_above_one = scalar.cmpf ogt, %hardswish_low, %c1_f32 : f32 + %hardswish_gate = scf.select %hardswish_above_one, %c1_f32, %hardswish_low : f32 + %hardswish = scalar.mulf %value, %hardswish_gate : f32 + %is_hardswish = index.cmp eq, %op, %op_hardswish : index + %hardswish_selected = scf.select %is_hardswish, %hardswish, %silu_selected : f32 + + %op_hardsigmoid = index.constant 17 : index + %hardsigmoid_bias = scalar.addf %value, %c3_f32 : f32 + %hardsigmoid_scaled = scalar.divf %hardsigmoid_bias, %c6_f32 : f32 + %hardsigmoid_above_zero = scalar.cmpf ogt, %hardsigmoid_scaled, %c0_f32 : f32 + %hardsigmoid_low = scf.select %hardsigmoid_above_zero, %hardsigmoid_scaled, %c0_f32 : f32 + %hardsigmoid_above_one = scalar.cmpf ogt, %hardsigmoid_low, %c1_f32 : f32 + %hardsigmoid = scf.select %hardsigmoid_above_one, %c1_f32, %hardsigmoid_low : f32 + %is_hardsigmoid = index.cmp eq, %op, %op_hardsigmoid : index + %hardsigmoid_selected = scf.select %is_hardsigmoid, %hardsigmoid, %hardswish_selected : f32 + + %op_exp = index.constant 18 : index + %exp_value = scalar.expf %value : f32 + %is_exp = index.cmp eq, %op, %op_exp : index + %exp_selected = scf.select %is_exp, %exp_value, %hardsigmoid_selected : f32 + + %op_expm1 = index.constant 19 : index + %expm1_exp = scalar.expf %value : f32 + %expm1 = scalar.subf %expm1_exp, %c1_f32 : f32 + %is_expm1 = index.cmp eq, %op, %op_expm1 : index + %expm1_selected = scf.select %is_expm1, %expm1, %exp_selected : f32 + + %op_gelu_erf = index.constant 20 : index + %gelu_erf_scaled = scalar.mulf %sqrt_2_inv, %value : f32 + %gelu_erf_value = scalar.erff %gelu_erf_scaled : f32 + %gelu_erf_one_plus = scalar.addf %c1_f32, %gelu_erf_value : f32 + %gelu_erf_half_x = scalar.mulf %c0_5_f32, %value : f32 + %gelu_erf = scalar.mulf %gelu_erf_half_x, %gelu_erf_one_plus : f32 + %is_gelu_erf = index.cmp eq, %op, %op_gelu_erf : index + %gelu_erf_selected = scf.select %is_gelu_erf, %gelu_erf, %expm1_selected : f32 + + %op_sqrt = index.constant 21 : index + %sqrt = scalar.sqrtf %value : f32 + %is_sqrt = index.cmp eq, %op, %op_sqrt : index + %sqrt_selected = scf.select %is_sqrt, %sqrt, %gelu_erf_selected : f32 + + %op_log = index.constant 22 : index + %log = scalar.logf %value : f32 + %is_log = index.cmp eq, %op, %op_log : index + %log_selected = scf.select %is_log, %log, %sqrt_selected : f32 + + %op_identity = index.constant 23 : index + %is_identity = index.cmp eq, %op, %op_identity : index + %result = scf.select %is_identity, %value, %log_selected : f32 + template.return %result : f32 +} + +template.def<@ggml.unary_f32.apply_vector4> device priority(1) @ggml_unary_f32_apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) { + %c0_f32 = vector.constant 0.0 : vector<4xf32> + %c0_5_f32 = vector.constant 0.5 : vector<4xf32> + %c1_f32 = vector.constant 1.0 : vector<4xf32> + %c2_f32 = vector.constant 2.0 : vector<4xf32> + %c3_f32 = vector.constant 3.0 : vector<4xf32> + %c6_f32 = vector.constant 6.0 : vector<4xf32> + %cn1_f32 = vector.constant -1.0 : vector<4xf32> + %cn2_f32 = vector.constant -2.0 : vector<4xf32> + %gelu_coef = vector.constant 0.044715 : vector<4xf32> + %gelu_quick_coef = vector.constant -1.702 : vector<4xf32> + %sqrt_2_over_pi = vector.constant 0.79788456080286544 : vector<4xf32> + %sqrt_2_inv = vector.constant 0.70710678118654745 : vector<4xf32> + + %op_neg = index.constant 0 : index + %negated = vector.negf %values : vector<4xf32> + %is_neg = index.cmp eq, %op, %op_neg : index + %neg_selected = scf.select %is_neg, %negated, %values : vector<4xf32> + + %op_abs = index.constant 1 : index + %abs_is_negative = vector.cmpf olt, %values, %c0_f32 : vector<4xf32> -> vector<4xi1> + %absolute = vector.select %abs_is_negative, %negated, %values : vector<4xf32> + %is_abs = index.cmp eq, %op, %op_abs : index + %abs_selected = scf.select %is_abs, %absolute, %neg_selected : vector<4xf32> + + %op_relu = index.constant 2 : index + %relu_is_positive = vector.cmpf ogt, %values, %c0_f32 : vector<4xf32> -> vector<4xi1> + %relu = vector.select %relu_is_positive, %values, %c0_f32 : vector<4xf32> + %is_relu = index.cmp eq, %op, %op_relu : index + %relu_selected = scf.select %is_relu, %relu, %abs_selected : vector<4xf32> + + %op_step = index.constant 3 : index + %step_is_positive = vector.cmpf ogt, %values, %c0_f32 : vector<4xf32> -> vector<4xi1> + %step = vector.select %step_is_positive, %c1_f32, %c0_f32 : vector<4xf32> + %is_step = index.cmp eq, %op, %op_step : index + %step_selected = scf.select %is_step, %step, %relu_selected : vector<4xf32> + + %op_sqr = index.constant 4 : index + %squared = vector.mulf %values, %values : vector<4xf32> + %is_sqr = index.cmp eq, %op, %op_sqr : index + %sqr_selected = scf.select %is_sqr, %squared, %step_selected : vector<4xf32> + + %op_sgn = index.constant 5 : index + %sgn_is_negative = vector.cmpf olt, %values, %c0_f32 : vector<4xf32> -> vector<4xi1> + %sgn_is_positive = vector.cmpf ogt, %values, %c0_f32 : vector<4xf32> -> vector<4xi1> + %negative_sgn = vector.select %sgn_is_negative, %cn1_f32, %c0_f32 : vector<4xf32> + %sgn = vector.select %sgn_is_positive, %c1_f32, %negative_sgn : vector<4xf32> + %is_sgn = index.cmp eq, %op, %op_sgn : index + %sgn_selected = scf.select %is_sgn, %sgn, %sqr_selected : vector<4xf32> + + %op_floor = index.constant 6 : index + %floored = vector.floorf %values : vector<4xf32> + %is_floor = index.cmp eq, %op, %op_floor : index + %floor_selected = scf.select %is_floor, %floored, %sgn_selected : vector<4xf32> + + %op_ceil = index.constant 7 : index + %ceiled = vector.ceilf %values : vector<4xf32> + %is_ceil = index.cmp eq, %op, %op_ceil : index + %ceil_selected = scf.select %is_ceil, %ceiled, %floor_selected : vector<4xf32> + + %op_round = index.constant 8 : index + %rounded = vector.roundf %values : vector<4xf32> + %is_round = index.cmp eq, %op, %op_round : index + %round_selected = scf.select %is_round, %rounded, %ceil_selected : vector<4xf32> + + %op_trunc = index.constant 9 : index + %truncated = vector.truncf %values : vector<4xf32> + %is_trunc = index.cmp eq, %op, %op_trunc : index + %trunc_selected = scf.select %is_trunc, %truncated, %round_selected : vector<4xf32> + + %op_tanh = index.constant 10 : index + %tanh_scaled = vector.mulf %cn2_f32, %values : vector<4xf32> + %tanh_exp = vector.expf %tanh_scaled : vector<4xf32> + %tanh_denominator = vector.addf %c1_f32, %tanh_exp : vector<4xf32> + %tanh_ratio = vector.divf %c2_f32, %tanh_denominator : vector<4xf32> + %tanh = vector.subf %tanh_ratio, %c1_f32 : vector<4xf32> + %is_tanh = index.cmp eq, %op, %op_tanh : index + %tanh_selected = scf.select %is_tanh, %tanh, %trunc_selected : vector<4xf32> + + %op_elu = index.constant 11 : index + %elu_is_positive = vector.cmpf ogt, %values, %c0_f32 : vector<4xf32> -> vector<4xi1> + %elu_exp = vector.expf %values : vector<4xf32> + %elu_expm1 = vector.subf %elu_exp, %c1_f32 : vector<4xf32> + %elu = vector.select %elu_is_positive, %values, %elu_expm1 : vector<4xf32> + %is_elu = index.cmp eq, %op, %op_elu : index + %elu_selected = scf.select %is_elu, %elu, %tanh_selected : vector<4xf32> + + %op_sigmoid = index.constant 12 : index + %sigmoid_negative = vector.negf %values : vector<4xf32> + %sigmoid_exp = vector.expf %sigmoid_negative : vector<4xf32> + %sigmoid_denominator = vector.addf %c1_f32, %sigmoid_exp : vector<4xf32> + %sigmoid = vector.divf %c1_f32, %sigmoid_denominator : vector<4xf32> + %is_sigmoid = index.cmp eq, %op, %op_sigmoid : index + %sigmoid_selected = scf.select %is_sigmoid, %sigmoid, %elu_selected : vector<4xf32> + + %op_gelu = index.constant 13 : index + %gelu_x2 = vector.mulf %values, %values : vector<4xf32> + %gelu_poly0 = vector.mulf %gelu_coef, %gelu_x2 : vector<4xf32> + %gelu_poly = vector.addf %c1_f32, %gelu_poly0 : vector<4xf32> + %gelu_inner0 = vector.mulf %values, %gelu_poly : vector<4xf32> + %gelu_inner = vector.mulf %sqrt_2_over_pi, %gelu_inner0 : vector<4xf32> + %gelu_tanh_scaled = vector.mulf %cn2_f32, %gelu_inner : vector<4xf32> + %gelu_tanh_exp = vector.expf %gelu_tanh_scaled : vector<4xf32> + %gelu_tanh_denominator = vector.addf %c1_f32, %gelu_tanh_exp : vector<4xf32> + %gelu_tanh_ratio = vector.divf %c2_f32, %gelu_tanh_denominator : vector<4xf32> + %gelu_tanh = vector.subf %gelu_tanh_ratio, %c1_f32 : vector<4xf32> + %gelu_one_plus = vector.addf %c1_f32, %gelu_tanh : vector<4xf32> + %gelu_half_x = vector.mulf %c0_5_f32, %values : vector<4xf32> + %gelu = vector.mulf %gelu_half_x, %gelu_one_plus : vector<4xf32> + %is_gelu = index.cmp eq, %op, %op_gelu : index + %gelu_selected = scf.select %is_gelu, %gelu, %sigmoid_selected : vector<4xf32> + + %op_gelu_quick = index.constant 14 : index + %gelu_quick_scaled = vector.mulf %gelu_quick_coef, %values : vector<4xf32> + %gelu_quick_exp = vector.expf %gelu_quick_scaled : vector<4xf32> + %gelu_quick_denominator = vector.addf %c1_f32, %gelu_quick_exp : vector<4xf32> + %gelu_quick_sigmoid = vector.divf %c1_f32, %gelu_quick_denominator : vector<4xf32> + %gelu_quick = vector.mulf %values, %gelu_quick_sigmoid : vector<4xf32> + %is_gelu_quick = index.cmp eq, %op, %op_gelu_quick : index + %gelu_quick_selected = scf.select %is_gelu_quick, %gelu_quick, %gelu_selected : vector<4xf32> + + %op_silu = index.constant 15 : index + %silu_negative = vector.negf %values : vector<4xf32> + %silu_exp = vector.expf %silu_negative : vector<4xf32> + %silu_denominator = vector.addf %c1_f32, %silu_exp : vector<4xf32> + %silu_sigmoid = vector.divf %c1_f32, %silu_denominator : vector<4xf32> + %silu = vector.mulf %values, %silu_sigmoid : vector<4xf32> + %is_silu = index.cmp eq, %op, %op_silu : index + %silu_selected = scf.select %is_silu, %silu, %gelu_quick_selected : vector<4xf32> + + %op_hardswish = index.constant 16 : index + %hardswish_bias = vector.addf %values, %c3_f32 : vector<4xf32> + %hardswish_scaled = vector.divf %hardswish_bias, %c6_f32 : vector<4xf32> + %hardswish_above_zero = vector.cmpf ogt, %hardswish_scaled, %c0_f32 : vector<4xf32> -> vector<4xi1> + %hardswish_low = vector.select %hardswish_above_zero, %hardswish_scaled, %c0_f32 : vector<4xf32> + %hardswish_above_one = vector.cmpf ogt, %hardswish_low, %c1_f32 : vector<4xf32> -> vector<4xi1> + %hardswish_gate = vector.select %hardswish_above_one, %c1_f32, %hardswish_low : vector<4xf32> + %hardswish = vector.mulf %values, %hardswish_gate : vector<4xf32> + %is_hardswish = index.cmp eq, %op, %op_hardswish : index + %hardswish_selected = scf.select %is_hardswish, %hardswish, %silu_selected : vector<4xf32> + + %op_hardsigmoid = index.constant 17 : index + %hardsigmoid_bias = vector.addf %values, %c3_f32 : vector<4xf32> + %hardsigmoid_scaled = vector.divf %hardsigmoid_bias, %c6_f32 : vector<4xf32> + %hardsigmoid_above_zero = vector.cmpf ogt, %hardsigmoid_scaled, %c0_f32 : vector<4xf32> -> vector<4xi1> + %hardsigmoid_low = vector.select %hardsigmoid_above_zero, %hardsigmoid_scaled, %c0_f32 : vector<4xf32> + %hardsigmoid_above_one = vector.cmpf ogt, %hardsigmoid_low, %c1_f32 : vector<4xf32> -> vector<4xi1> + %hardsigmoid = vector.select %hardsigmoid_above_one, %c1_f32, %hardsigmoid_low : vector<4xf32> + %is_hardsigmoid = index.cmp eq, %op, %op_hardsigmoid : index + %hardsigmoid_selected = scf.select %is_hardsigmoid, %hardsigmoid, %hardswish_selected : vector<4xf32> + + %op_exp = index.constant 18 : index + %exp_value = vector.expf %values : vector<4xf32> + %is_exp = index.cmp eq, %op, %op_exp : index + %exp_selected = scf.select %is_exp, %exp_value, %hardsigmoid_selected : vector<4xf32> + + %op_expm1 = index.constant 19 : index + %expm1_exp = vector.expf %values : vector<4xf32> + %expm1 = vector.subf %expm1_exp, %c1_f32 : vector<4xf32> + %is_expm1 = index.cmp eq, %op, %op_expm1 : index + %expm1_selected = scf.select %is_expm1, %expm1, %exp_selected : vector<4xf32> + + %op_gelu_erf = index.constant 20 : index + %gelu_erf_scaled = vector.mulf %sqrt_2_inv, %values : vector<4xf32> + %gelu_erf_value = vector.erff %gelu_erf_scaled : vector<4xf32> + %gelu_erf_one_plus = vector.addf %c1_f32, %gelu_erf_value : vector<4xf32> + %gelu_erf_half_x = vector.mulf %c0_5_f32, %values : vector<4xf32> + %gelu_erf = vector.mulf %gelu_erf_half_x, %gelu_erf_one_plus : vector<4xf32> + %is_gelu_erf = index.cmp eq, %op, %op_gelu_erf : index + %gelu_erf_selected = scf.select %is_gelu_erf, %gelu_erf, %expm1_selected : vector<4xf32> + + %op_sqrt = index.constant 21 : index + %sqrt = vector.sqrtf %values : vector<4xf32> + %is_sqrt = index.cmp eq, %op, %op_sqrt : index + %sqrt_selected = scf.select %is_sqrt, %sqrt, %gelu_erf_selected : vector<4xf32> + + %op_log = index.constant 22 : index + %log = vector.logf %values : vector<4xf32> + %is_log = index.cmp eq, %op, %op_log : index + %log_selected = scf.select %is_log, %log, %sqrt_selected : vector<4xf32> + + %op_identity = index.constant 23 : index + %is_identity = index.cmp eq, %op, %op_identity : index + %result = scf.select %is_identity, %values, %log_selected : vector<4xf32> + template.return %result : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/f16_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/f16_f16.loom new file mode 100644 index 000000000000..e64b41ef147a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/f16_f16.loom @@ -0,0 +1,10 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Loads four adjacent F16 weights for FP16 matrix staging. +func.def inline @ggml_f16_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xf16> + %values = vector.load %row_view[%k] : view<[%bounded_input_size]xf16> -> vector<4xf16> + func.return %values : vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/f32_f16.loom new file mode 100644 index 000000000000..1de466245d5f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/f32_f16.loom @@ -0,0 +1,11 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Loads four adjacent F32 weights and truncates them for FP16 matrix staging. +func.def inline @ggml_f32_f16_vector4(%weight: buffer, %row_byte_base: offset, %input_size: index, %k: index) -> (vector<4xf16>) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %row_view = buffer.view %weight[%row_byte_base] : buffer -> view<[%bounded_input_size]xf32> + %values_f32 = vector.load %row_view[%k] : view<[%bounded_input_size]xf32> -> vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/completion_counter_reduce.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/completion_counter_reduce.loom new file mode 100644 index 000000000000..aef47951d5a7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/completion_counter_reduce.loom @@ -0,0 +1,108 @@ +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +template.decl @ggml.flash_attention.decode_split.pack_completed_q8(%key_value_head: index, %output: buffer, %q8_output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_completed.cooperative(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_completed.direct(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_fused(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) + +// Completes bounded decode contexts inside the last arriving producer +// workgroup. Four KV-head workgroups reduce their own query heads while all +// other producers retire, erasing a second dispatch and its execution barrier. +// Capacity selects the algorithm and partial layout; producer count owns the +// issue-time completion threshold and active reduction prefix. +template.def<@ggml.flash_attention.decode_split.reduce_fused> device priority(20) @ggml_flash_attention_decode_split_reduce_fused_direct_f32(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) where [range(%key_value_token_capacity, 64, 256)] { + %partial_block_capacity, %producer_block_count = index.assume %partial_block_capacity0, %producer_block_count0 [range(%partial_block_capacity0, 1, 4), range(%producer_block_count0, 1, 4), le(%producer_block_count0, %partial_block_capacity0)] : index, index + %key_value_head_count = config.get @ggml.flash_attention.key_value_head_count : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %completion_counter_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %completion_counter, %output : buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + // Publish every producer's partial stores before the leader advances one + // workgroup arrival. The last arrival then acquires every partial. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%key_value_head] {ordering = acq_rel, scope = device} : i32, view<[%key_value_head_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %key_value_block_count_i32 = index.cast %producer_block_count : index to i32 + %last_block_ordinal_i32 = scalar.subi %key_value_block_count_i32, %c1_i32 : i32 + %negative_key_value_block_count_i32 = scalar.subi %c0_i32, %key_value_block_count_i32 : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %last_block_ordinal_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + template.apply<@ggml.flash_attention.decode_split.reduce_completed.direct>(%partial_block_capacity, %producer_block_count, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias) : (index, index, buffer, buffer, buffer, buffer) + scf.if %publish_q8 { + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.flash_attention.decode_split.pack_completed_q8>(%key_value_head, %output_noalias, %q8_output) : (index, buffer, buffer) + } + // Do not expose the reset until every final F32 and Q8 store completes. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_key_value_block_count_i32, %completion_counter_view[%key_value_head] {ordering = release, scope = device} : i32, view<[%key_value_head_count]xi32> + } + } + template.return +} +// Contexts without a proven short bound use the cooperative completion path. +template.def<@ggml.flash_attention.decode_split.reduce_fused> device priority(10) @ggml_flash_attention_decode_split_reduce_fused_cooperative_f32(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) where [range(%key_value_token_capacity, 257, 2048)] { + %partial_block_capacity, %producer_block_count = index.assume %partial_block_capacity0, %producer_block_count0 [range(%partial_block_capacity0, 1, 32), range(%producer_block_count0, 1, 32), le(%producer_block_count0, %partial_block_capacity0)] : index, index + %key_value_head_count = config.get @ggml.flash_attention.key_value_head_count : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %completion_counter_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %completion_counter, %output : buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + // Publish every producer's partial stores before the leader advances one + // workgroup arrival. The last arrival then acquires every partial. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%key_value_head] {ordering = acq_rel, scope = device} : i32, view<[%key_value_head_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %key_value_block_count_i32 = index.cast %producer_block_count : index to i32 + %last_block_ordinal_i32 = scalar.subi %key_value_block_count_i32, %c1_i32 : i32 + %negative_key_value_block_count_i32 = scalar.subi %c0_i32, %key_value_block_count_i32 : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %last_block_ordinal_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + template.apply<@ggml.flash_attention.decode_split.reduce_completed.cooperative>(%partial_block_capacity, %producer_block_count, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias) : (index, index, buffer, buffer, buffer, buffer) + scf.if %publish_q8 { + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.flash_attention.decode_split.pack_completed_q8>(%key_value_head, %output_noalias, %q8_output) : (index, buffer, buffer) + } + // Do not expose the reset until every final F32 and Q8 store completes. + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_key_value_block_count_i32, %completion_counter_view[%key_value_head] {ordering = release, scope = device} : i32, view<[%key_value_head_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/decode_partials.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/decode_partials.loom new file mode 100644 index 000000000000..36d2ecca1461 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/decode_partials.loom @@ -0,0 +1,516 @@ +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.qk_head_size : %value: index where [range(%value, 16, 512), mul(%value, 16)] + +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] + +config.decl @ggml.flash_attention.attention_scale : f32 + +config.decl @ggml.flash_attention.decode.key_value_token_capacity : %value: index where [range(%value, 64, 32768)] + +template.decl @ggml.flash_attention.decode_split.produce_partials(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) + +template.decl @ggml.flash_attention.decode_split.produce_partials.active(%key_value_token_count: index, %partial_block_capacity0: index, %launched_block_count0: index, %query: buffer, %key: buffer, %value: buffer, %lane_mask: f32, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) + +// Computes one active online-softmax partial for a 64-row KV block. Both the +// fused short-context export and the two-dispatch long-context export reach +// this body through the block-classifying producer below, keeping their +// different binding contracts honest without duplicating the attention math. +// Partial storage retains its capacity-specialized block-axis stride while the +// launched block count bounds issue-time work. Keeping those values distinct +// preserves constant address arithmetic across changing visible prefixes. +template.def<@ggml.flash_attention.decode_split.produce_partials.active> device @ggml_flash_attention_decode_split_produce_active_partials_body_f32_f16_wmma(%key_value_token_count: index, %partial_block_capacity0: index, %launched_block_count0: index, %query: buffer, %key: buffer, %value: buffer, %lane_mask: f32, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { + %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %partial_block_capacity, %launched_block_count = index.assume %partial_block_capacity0, %launched_block_count0 [range(%partial_block_capacity0, 1, 512), range(%launched_block_count0, 1, 512), le(%launched_block_count0, %partial_block_capacity0)] : index, 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 + %qk_head_size0 = config.get @ggml.flash_attention.qk_head_size : index + %qk_head_size = index.assume %qk_head_size0 [range(%qk_head_size0, 16, 512), mul(%qk_head_size0, 16)] : index + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %workgroup_x0 = kernel.workgroup.id : index + %workgroup_y0 = kernel.workgroup.id : index + %workgroup_x_in_launch = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 511)] : index + %workgroup_y = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : 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 + %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 + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %f16_bytes = index.constant 2 : offset + %score_stage_bytes = index.constant 6144 : offset + %probability_stage_bytes = index.constant 3072 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %attention_scale = config.get @ggml.flash_attention.attention_scale : f32 + %c0_f16 = scalar.constant 0.0 : f16 + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f16x16 = vector.constant 0.0 : vector<16xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %workgroup_x, %launch_partial_block_capacity, %launch_launched_block_count = index.assume %workgroup_x_in_launch, %partial_block_capacity, %launched_block_count [lt(%workgroup_x_in_launch, %launched_block_count)] : index, index, index + %tail_key_value_token_count = index.rem %bounded_key_value_token_count, %c64 : index + %has_no_tail = index.cmp eq, %tail_key_value_token_count, %c0 : index + %active_padded_key_value_token_count = index.add %bounded_key_value_token_count, %c63 : index + %active_key_value_block_count = index.div %active_padded_key_value_token_count, %c64 : index + %last_block_ordinal = index.sub %active_key_value_block_count, %c1 : index + %block_ordinal = index.add %workgroup_x, %c0 : index + %is_not_last_block = index.cmp ne, %block_ordinal, %last_block_ordinal : index + %is_full_block = scalar.ori %has_no_tail, %is_not_last_block : i1 + %key_value_head, %launch_key_value_head_count = index.assume %workgroup_y, %key_value_head_count [lt(%workgroup_y, %key_value_head_count)] : index, index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %key_width = index.mul %key_value_head_count, %qk_head_size : index + %value_width = index.mul %key_value_head_count, %value_head_size : index + %key_head_base = index.mul %key_value_head, %qk_head_size : index + %value_head_base = index.mul %key_value_head, %value_head_size : index + %padded_value_head_size = index.add %value_head_size, %c127 : index + %output_tile_count = index.div %padded_value_head_size, %c128 : index + %output_stage_size = index.mul %output_tile_count, %c128 : index + %query_stage_stride = index.add %qk_head_size, %c8 : index + %query_stage_element_count = index.mul %c16, %query_stage_stride : index + %query_stage_bytes = index.scale %query_stage_element_count, %f16_bytes : index, offset -> offset + %query_element_count = index.mul %c16, %qk_head_size : index + %query_load_iteration_count = index.div %query_element_count, %c256 : index + %product_stage_element_count = index.mul %c16, %output_stage_size : index + %product_stage_bytes = index.scale %product_stage_element_count, %f16_bytes : index, offset -> offset + %key_origin = index.mul %block_ordinal, %c64 : index + %subgroup_score_column = index.mul %subgroup, %c16 : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head_base, %query_row1 : index + %query_head2 = index.add %query_head_base, %query_row2 : index + %query_head3 = index.add %query_head_base, %query_row3 : index + %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index + %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + %query_head_valid2 = index.cmp ult, %query_head2, %query_head_count : index + %query_head_valid3 = index.cmp ult, %query_head3, %query_head_count : index + %query_valid = vector.from_elements %query_head_valid0, %query_head_valid1, %query_head_valid2, %query_head_valid3 : vector<4xi1> + %lane_has_output = index.cmp ult, %lane, %c32 : index + %lane_is_zero = index.cmp eq, %lane, %c0 : index + %query_transposed_layout = encoding.layout.strided [1, %query_stage_stride] : encoding + %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding + %query_noalias, %key_noalias, %value_noalias, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias = buffer.assume.noalias %query, %key, %value, %partial_max, %partial_sum, %partial_output : buffer, buffer, buffer, buffer, buffer, buffer + %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer + %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer + %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer + %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer + %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x[%qk_head_size]xf32> + %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> + %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16> + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x[%value_head_size]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 + %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %tail_value_stage = buffer.alloca align(16) %product_stage_bytes : buffer + %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x[%query_stage_stride]xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x[%output_stage_size]xf16> + %tail_key_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<64x16xf16> + %tail_value_stage_view = buffer.view %tail_value_stage[%c0_offset] : buffer -> view<16x[%output_stage_size]xf16> + scf.for %load_iteration = [%c0 to %query_load_iteration_count step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %local_query_row = index.div %linear, %qk_head_size : index + %query_channel = index.rem %linear, %qk_head_size : index + %local_query_head = index.add %query_head_base, %local_query_row : index + %local_query_head_valid = index.cmp ult, %local_query_head, %query_head_count : index + %query_value = scf.if %local_query_head_valid -> (f16) { + %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x[%qk_head_size]xf32> -> f32 + %scaled = scalar.mulf %loaded, %attention_scale : f32 + %truncated = scalar.fptrunc %scaled : f32 to f16 + scf.yield %truncated : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x[%query_stage_stride]xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %score_init_values = vector.constant 0.0 : vector<4xf32> + %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> + %score_fragment = scf.if %is_full_block -> (vector<4xf32>) { + %full_score_fragment = scf.for %head_tile = [%c0 to %qk_head_size step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + // Bound this view by the proven tile end so the full vector footprint is + // visible without treating physical tail padding as logical storage. + %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index + %score_key_end0 = index.add %score_key_origin0, %c16 : index + %score_key_origin, %score_key_end = index.assume %score_key_origin0, %score_key_end0 [le(%score_key_end0, %bounded_key_value_token_count)] : index, index + %full_key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%score_key_end]x[%key_width]xf16> + %key_channel = index.add %key_head_base, %head_tile : index + %key_fragment = vector.fragment.load %full_key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%score_key_end]x[%key_width]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<[%qk_head_size]x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + scf.yield %full_score_fragment : vector<4xf32> + } else { + // Stage one 64x16 K panel at a time. Every physical load is guarded, so + // callers need no initialized padding beyond the logical KV length. + %tail_score_fragment = scf.for %head_tile = [%c0 to %qk_head_size step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) { + scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %key_row = index.div %linear, %c16 : index + %head_channel = index.rem %linear, %c16 : index + %key_token = index.add %key_origin, %key_row : index + %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index + %key_value = scf.if %key_valid -> (f16) { + %head_channel0 = index.add %head_tile, %head_channel : index + %key_channel = index.add %key_head_base, %head_channel0 : index + %loaded = view.load %key_view[%key_token, %key_channel] : view<[%bounded_key_value_token_count]x[%key_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %key_value, %tail_key_stage_view[%key_row, %head_channel] : f16, view<64x16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %key_fragment = vector.fragment.load %tail_key_stage_view[%subgroup_score_column, %c0] shape [%m, %k] : view<64x16xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<[%qk_head_size]x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_score_accumulator : vector<4xf32> + } + scf.yield %tail_score_fragment : vector<4xf32> + } + vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %key_token = index.add %key_origin, %lane : index + %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 + %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 + %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 + %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 + %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index + %score_valid0 = scalar.andi %query_head_valid0, %key_valid : i1 + %score_valid1 = scalar.andi %query_head_valid1, %key_valid : i1 + %score_valid2 = scalar.andi %query_head_valid2, %key_valid : i1 + %score_valid3 = scalar.andi %query_head_valid3, %key_valid : i1 + %masked_score0 = scf.if %score_valid0 -> (f32) { + %score = scalar.addf %raw_score0, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score1 = scf.if %score_valid1 -> (f32) { + %score = scalar.addf %raw_score1, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score2 = scf.if %score_valid2 -> (f32) { + %score = scalar.addf %raw_score2, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score3 = scf.if %score_valid3 -> (f32) { + %score = scalar.addf %raw_score3, %lane_mask : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %score_delta = vector.subf %masked_scores, %block_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %output_tile = [%c0 to %output_tile_count step %c1] unroll { + %output_tile_channel = index.mul %output_tile, %c128 : index + %subgroup_output_base = index.mul %subgroup, %c32 : index + %subgroup_output_channel0 = index.add %output_tile_channel, %subgroup_output_base : index + %subgroup_output_channel = index.assume %subgroup_output_channel0 [range(%subgroup_output_channel0, 0, 480)] : index + %subgroup_output_channel1_0 = index.add %subgroup_output_channel, %c16 : index + %subgroup_output_channel1 = index.assume %subgroup_output_channel1_0 [range(%subgroup_output_channel1_0, 16, 496)] : index + %subgroup_output_channel_valid = index.cmp ult, %subgroup_output_channel, %value_head_size : index + %subgroup_output_channel1_valid = index.cmp ult, %subgroup_output_channel1, %value_head_size : index + %lane_output_base = index.mul %lane, %c4 : index + %lane_output_channel0 = index.add %output_tile_channel, %lane_output_base : index + %lane_output_channel = index.assume %lane_output_channel0 [range(%lane_output_channel0, 0, 636)] : index + %lane_output_channel_valid = index.cmp ult, %lane_output_channel, %value_head_size : index + %lane_publishes_output = scalar.andi %lane_has_output, %lane_output_channel_valid : i1 + %safe_lane_output_channel0 = scf.select %lane_output_channel_valid, %lane_output_channel, %c0 : index + %safe_lane_output_channel_end0 = index.add %safe_lane_output_channel0, %c4 : index + %safe_lane_output_channel, %safe_lane_output_channel_end = index.assume %safe_lane_output_channel0, %safe_lane_output_channel_end0 [le(%safe_lane_output_channel_end0, %value_head_size), le(%safe_lane_output_channel_end0, %output_stage_size)] : index, index + %product_init0 = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %product_init1 = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %safe_subgroup_output_channel0 = scf.select %subgroup_output_channel_valid, %subgroup_output_channel, %c0 : index + %safe_subgroup_output_channel1_0 = scf.select %subgroup_output_channel1_valid, %subgroup_output_channel1, %c0 : index + %safe_subgroup_output_channel_end0 = index.add %safe_subgroup_output_channel0, %c16 : index + %safe_subgroup_output_channel1_end0 = index.add %safe_subgroup_output_channel1_0, %c16 : index + %safe_subgroup_output_channel, %safe_subgroup_output_channel_end = index.assume %safe_subgroup_output_channel0, %safe_subgroup_output_channel_end0 [le(%safe_subgroup_output_channel_end0, %value_head_size), le(%safe_subgroup_output_channel_end0, %output_stage_size)] : index, index + %safe_subgroup_output_channel1, %safe_subgroup_output_channel1_end = index.assume %safe_subgroup_output_channel1_0, %safe_subgroup_output_channel1_end0 [le(%safe_subgroup_output_channel1_end0, %value_head_size), le(%safe_subgroup_output_channel1_end0, %output_stage_size)] : index, index + %value_channel0 = index.add %value_head_base, %safe_subgroup_output_channel : index + %value_channel1 = index.add %value_head_base, %safe_subgroup_output_channel1 : index + %value_channel0_end0 = index.add %value_head_base, %safe_subgroup_output_channel_end : index + %value_channel1_end0 = index.add %value_head_base, %safe_subgroup_output_channel1_end : index + %safe_value_channel0, %safe_value_channel0_end = index.assume %value_channel0, %value_channel0_end0 [le(%value_channel0_end0, %value_width)] : index, index + %safe_value_channel1, %safe_value_channel1_end = index.assume %value_channel1, %value_channel1_end0 [le(%value_channel1_end0, %value_width)] : index, index + %last_value_fragment_channel = index.sub %value_width, %c16 : index + %clamped_value_channel0 = index.min %safe_value_channel0, %last_value_fragment_channel : index + %clamped_value_channel1 = index.min %safe_value_channel1, %last_value_fragment_channel : index + %last_stage_fragment_channel = index.sub %output_stage_size, %c16 : index + %clamped_stage_channel0 = index.min %safe_subgroup_output_channel, %last_stage_fragment_channel : index + %clamped_stage_channel1 = index.min %safe_subgroup_output_channel1, %last_stage_fragment_channel : index + %product_fragment0, %product_fragment1 = scf.if %is_full_block -> (vector<8xf16>, vector<8xf16>) { + %full_product_fragment0, %full_product_fragment1 = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator0 = %product_init0 : vector<8xf16>, %product_accumulator1 = %product_init1 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>) unroll { + %value_token0 = index.add %key_origin, %key_tile : index + %value_token_end0 = index.add %value_token0, %c16 : index + %value_token, %value_token_end = index.assume %value_token0, %value_token_end0 [le(%value_token_end0, %bounded_key_value_token_count)] : index, index + %full_value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%value_token_end]x[%value_width]xf16> + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + %value_fragment0 = vector.fragment.load %full_value_view[%value_token, %clamped_value_channel0] shape [%k, %n] : view<[%value_token_end]x[%value_width]xf16> -> vector<16xf16> + %value_fragment1 = vector.fragment.load %full_value_view[%value_token, %clamped_value_channel1] shape [%k, %n] : view<[%value_token_end]x[%value_width]xf16> -> vector<16xf16> + %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> + } + scf.yield %full_product_fragment0, %full_product_fragment1 : vector<8xf16>, vector<8xf16> + } else { + %tail_product_fragment0, %tail_product_fragment1 = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator0 = %product_init0 : vector<8xf16>, %product_accumulator1 = %product_init1 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>) unroll { + scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %key_row = index.div %linear, %c128 : index + %tile_value_channel = index.rem %linear, %c128 : index + %value_channel = index.add %output_tile_channel, %tile_value_channel : index + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.add %value_token0, %key_row : index + %value_token_valid = index.cmp ult, %value_token, %bounded_key_value_token_count : index + %value_channel_valid = index.cmp ult, %value_channel, %value_head_size : index + %value_valid = scalar.andi %value_token_valid, %value_channel_valid : i1 + %value_element = scf.if %value_valid -> (f16) { + %global_value_channel = index.add %value_head_base, %value_channel : index + %loaded = view.load %value_view[%value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%value_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %value_element, %tail_value_stage_view[%key_row, %value_channel] : f16, view<16x[%output_stage_size]xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> + // Bound each fragment without changing the physical LDS row stride. + %value_fragment0 = vector.fragment.load %tail_value_stage_view[%c0, %clamped_stage_channel0] shape [%k, %n] : view<16x[%output_stage_size]xf16> -> vector<16xf16> + %value_fragment1 = vector.fragment.load %tail_value_stage_view[%c0, %clamped_stage_channel1] shape [%k, %n] : view<16x[%output_stage_size]xf16> -> vector<16xf16> + %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> + } + scf.yield %tail_product_fragment0, %tail_product_fragment1 : vector<8xf16>, vector<8xf16> + } + vector.fragment.store %product_fragment0, %product_stage_view[%c0, %subgroup_output_channel] shape [%m, %n] : vector<8xf16>, view<16x[%output_stage_size]xf16> + vector.fragment.store %product_fragment1, %product_stage_view[%c0, %subgroup_output_channel1] shape [%m, %n] : vector<8xf16>, view<16x[%output_stage_size]xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %lane_publishes_output { + %block_output0 = vector.load %product_stage_view[%query_row0, %safe_lane_output_channel] : view<16x[%output_stage_size]xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %safe_lane_output_channel] : view<16x[%output_stage_size]xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %safe_lane_output_channel] : view<16x[%output_stage_size]xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %safe_lane_output_channel] : view<16x[%output_stage_size]xf16> -> vector<4xf16> + scf.if %query_head_valid0 { + vector.store %block_output0, %partial_output_view[%key_value_head, %block_ordinal, %query_row0, %safe_lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x[%value_head_size]xf16> + } + scf.if %query_head_valid1 { + vector.store %block_output1, %partial_output_view[%key_value_head, %block_ordinal, %query_row1, %safe_lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x[%value_head_size]xf16> + } + scf.if %query_head_valid2 { + vector.store %block_output2, %partial_output_view[%key_value_head, %block_ordinal, %query_row2, %safe_lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x[%value_head_size]xf16> + } + scf.if %query_head_valid3 { + vector.store %block_output3, %partial_output_view[%key_value_head, %block_ordinal, %query_row3, %safe_lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x[%value_head_size]xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %lane_is_zero { + scf.if %query_head_valid0 { + %maximum = vector.extract %block_max[0] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[0] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + scf.if %query_head_valid1 { + %maximum = vector.extract %block_max[1] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[1] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + scf.if %query_head_valid2 { + %maximum = vector.extract %block_max[2] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[2] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + scf.if %query_head_valid3 { + %maximum = vector.extract %block_max[3] : vector<4xf32> -> f32 + %sum = vector.extract %block_sum[3] : vector<4xf32> -> f32 + view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> + } + } + template.return +} + +// Fully masked splits publish the online-softmax identity without running QK +// or P*V. Valid additive masks contain finite F16 values or negative infinity, +// and every finite F16 value is greater than -1e30. Comparing after extension +// therefore distinguishes active rows from masked rows exactly. +template.def<@ggml.flash_attention.decode_split.produce_partials> device @ggml_flash_attention_decode_split_produce_partials_body_f32_f16_wmma(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { + %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 + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %workgroup_x0 = kernel.workgroup.id : index + %workgroup_y0 = kernel.workgroup.id : index + %workgroup_x_in_launch = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 511)] : index + %workgroup_y = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : 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 + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c63 = index.constant 63 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count0 = index.div %padded_key_value_token_capacity, %c64 : index + %key_value_block_count = index.assume %key_value_block_count0 [range(%key_value_block_count0, 1, 512)] : index + %workgroup_x, %launch_key_value_block_count = index.assume %workgroup_x_in_launch, %key_value_block_count [lt(%workgroup_x_in_launch, %key_value_block_count)] : index, index + %block_ordinal = index.add %workgroup_x, %c0 : index + %key_origin = index.mul %block_ordinal, %c64 : index + %key_token = index.add %key_origin, %lane : index + %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index + %mask_aligned = buffer.assume.alignment %mask {minimum_alignment = 16} : buffer + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]xf16> + %lane_mask = scf.if %key_valid -> (f32) { + %mask_f16 = view.load %mask_view[%key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + scf.yield %mask_f32 : f32 + } else { + scf.yield %negative_large : f32 + } + %block_mask_maximum = kernel.workgroup.reduce %lane_mask : f32 + %block_has_attention = scalar.cmpf ogt, %block_mask_maximum, %negative_large : f32 + scf.if %block_has_attention { + template.apply<@ggml.flash_attention.decode_split.produce_partials.active>(%bounded_key_value_token_count, %launch_key_value_block_count, %launch_key_value_block_count, %query, %key, %value, %lane_mask, %partial_max, %partial_sum, %partial_output) : (index, index, index, buffer, buffer, buffer, f32, buffer, buffer, buffer) + } else { + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %key_value_head = index.add %workgroup_y, %c0 : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head_base, %query_row1 : index + %query_head2 = index.add %query_head_base, %query_row2 : index + %query_head3 = index.add %query_head_base, %query_row3 : index + %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index + %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index + %query_head_valid2 = index.cmp ult, %query_head2, %query_head_count : index + %query_head_valid3 = index.cmp ult, %query_head3, %query_head_count : index + %padded_value_head_size = index.add %value_head_size, %c127 : index + %output_tile_count = index.div %padded_value_head_size, %c128 : index + %lane_output_base = index.mul %lane, %c4 : index + %lane_has_output = index.cmp ult, %lane, %c32 : index + %lane_is_zero = index.cmp eq, %lane, %c0 : index + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output : buffer, buffer, buffer + %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> + scf.for %output_tile = [%c0 to %output_tile_count step %c1] unroll { + %output_tile_channel = index.mul %output_tile, %c128 : index + %lane_output_channel = index.add %output_tile_channel, %lane_output_base : index + %lane_output_channel_valid = index.cmp ult, %lane_output_channel, %value_head_size : index + %lane_publishes_output = scalar.andi %lane_has_output, %lane_output_channel_valid : i1 + scf.if %lane_publishes_output { + scf.if %query_head_valid0 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row0, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> + } + scf.if %query_head_valid1 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row1, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> + } + scf.if %query_head_valid2 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row2, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> + } + scf.if %query_head_valid3 { + vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row3, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> + } + } + } + scf.if %lane_is_zero { + scf.if %query_head_valid0 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + scf.if %query_head_valid1 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + scf.if %query_head_valid2 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + scf.if %query_head_valid3 { + view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/decode_reduce.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/decode_reduce.loom new file mode 100644 index 000000000000..97438fe679ef --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/decode_reduce.loom @@ -0,0 +1,205 @@ +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] + +template.decl @ggml.flash_attention.decode_split.reduce_completed.cooperative(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_completed.direct(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.publish_f32.scalar(%query_head_count: index, %value_head_size: index, %query_head: index, %output_channel: index, %sum: f32, %unnormalized_output: f32, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.publish_f32.vector2(%query_head_count: index, %value_head_size: index, %query_head_valid0: i1, %query_head_valid1: i1, %lane_output_channel_valid: i1, %query_head0: index, %query_head1: index, %lane_output_channel: index, %sum0: f32, %sum1: f32, %unnormalized_output0: vector<2xf32>, %unnormalized_output1: vector<2xf32>, %output: buffer) + +// Up to four split-K blocks are cheapest to fold directly. Each workitem owns +// one output element, so the reducer needs no LDS or subgroup synchronization. +template.def<@ggml.flash_attention.decode_split.reduce_completed.direct> device @ggml_flash_attention_decode_split_reduce_completed_direct_f32(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) where [range(%partial_block_capacity0, 1, 4)] { + %partial_block_capacity, %active_block_count = index.assume %partial_block_capacity0, %active_block_count0 [range(%partial_block_capacity0, 1, 4), range(%active_block_count0, 1, 4), le(%active_block_count0, %partial_block_capacity0)] : index, 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 + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %partial_max_aligned = buffer.assume.alignment %partial_max {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> + %partial_element_count = index.mul %query_heads_per_key_value_head, %value_head_size : index + scf.for %linear = [%workitem to %partial_element_count step %c256] { + %query_row = index.div %linear, %value_head_size : index + %output_channel = index.rem %linear, %value_head_size : index + %query_head = index.add %query_head_base, %query_row : index + %query_head_valid = index.cmp ult, %query_head, %query_head_count : index + scf.if %query_head_valid { + %maximum = scf.for %block = [%c0 to %active_block_count step %c1](%running_maximum = %negative_large : f32) -> (f32) unroll schedule(interleaved) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %next_maximum = scalar.maxnumf %running_maximum, %block_maximum : f32 + scf.yield %next_maximum : f32 + } + %sum, %unnormalized_output = scf.for %block = [%c0 to %active_block_count step %c1](%running_sum = %c0_f32 : f32, %running_output = %c0_f32 : f32) -> (f32, f32) unroll schedule(interleaved) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum : f32 + %scale = scalar.expf %delta : f32 + %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + %next_sum = scalar.addf %running_sum, %scaled_sum : f32 + %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> -> f16 + %block_output = scalar.extf %block_output_f16 : f16 to f32 + %scaled_output = scalar.mulf %block_output, %scale : f32 + %next_output = scalar.addf %running_output, %scaled_output : f32 + scf.yield %next_sum, %next_output : f32, f32 + } + template.apply<@ggml.flash_attention.decode_split.publish_f32.scalar>(%query_head_count, %value_head_size, %query_head, %output_channel, %sum, %unnormalized_output, %output_aligned) : (index, index, index, index, f32, f32, buffer) + } + } + template.return +} + +// Longer bounded contexts amortize a cooperative reducer. Each wave folds two +// query rows, stages the per-block scales in LDS, and writes two output channels +// per lane. +template.def<@ggml.flash_attention.decode_split.reduce_completed.cooperative> device @ggml_flash_attention_decode_split_reduce_completed_cooperative_f32(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) { + %partial_block_capacity, %active_block_count = index.assume %partial_block_capacity0, %active_block_count0 [range(%partial_block_capacity0, 1, 32), range(%active_block_count0, 1, 32), le(%active_block_count0, %partial_block_capacity0)] : index, 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 + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %workgroup_y0 = kernel.workgroup.id : index + %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : 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 + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %scale_stage_bytes = index.constant 1024 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %c0_f16x2 = vector.constant 0.0 : vector<2xf16> + %c0_f32x2 = vector.constant 0.0 : vector<2xf32> + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %padded_value_head_size = index.add %value_head_size, %c127 : index + %output_tile_count = index.div %padded_value_head_size, %c128 : index + %lane_output_base = index.mul %lane, %c2 : index + %lane_has_block = index.cmp ult, %lane, %active_block_count : index + %partial_max_aligned = buffer.assume.alignment %partial_max {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<4x32x2xf32> + %padded_reduction_row_count = index.add %query_heads_per_key_value_head, %c7 : index + %reduction_phase_count = index.div %padded_reduction_row_count, %c8 : index + scf.for %phase = [%c0 to %reduction_phase_count step %c1] unroll { + %phase_row_base = index.mul %phase, %c8 : index + %subgroup_row_base = index.mul %subgroup, %c2 : index + %query_row0 = index.add %phase_row_base, %subgroup_row_base : index + %query_row1 = index.add %query_row0, %c1 : index + %query_head0 = index.add %query_head_base, %query_row0 : index + %query_head1 = index.add %query_head0, %c1 : index + // Inactive rows must not write query heads owned by the next KV head. + %query_head_valid0 = index.cmp ult, %query_row0, %query_heads_per_key_value_head : index + %query_head_valid1 = index.cmp ult, %query_row1, %query_heads_per_key_value_head : index + %safe_query_row0 = scf.select %query_head_valid0, %query_row0, %c0 : index + %safe_query_row1 = scf.select %query_head_valid1, %query_row1, %c0 : index + %reducer_active0 = scalar.andi %lane_has_block, %query_head_valid0 : i1 + %reducer_active1 = scalar.andi %lane_has_block, %query_head_valid1 : i1 + %local_maximum0 = scf.if %reducer_active0 -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + scf.yield %block_maximum : f32 + } else { + scf.yield %negative_large : f32 + } + %local_maximum1 = scf.if %reducer_active1 -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + scf.yield %block_maximum : f32 + } else { + scf.yield %negative_large : f32 + } + %local_maximums = vector.from_elements %local_maximum0, %local_maximum1 : vector<2xf32> + %maximums = kernel.subgroup.reduce %local_maximums : vector<2xf32> + %maximum0 = vector.extract %maximums[0] : vector<2xf32> -> f32 + %maximum1 = vector.extract %maximums[1] : vector<2xf32> -> f32 + %local_sum0, %scale0 = scf.if %reducer_active0 -> (f32, f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum0 : f32 + %scale = scalar.expf %delta : f32 + %block_sum = view.load %partial_sum_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + scf.yield %scaled_sum, %scale : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + %local_sum1, %scale1 = scf.if %reducer_active1 -> (f32, f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum1 : f32 + %scale = scalar.expf %delta : f32 + %block_sum = view.load %partial_sum_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + scf.yield %scaled_sum, %scale : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + scf.if %lane_has_block { + %bounded_block_lane = index.assume %lane [range(%lane, 0, 31)] : index + %scales = vector.from_elements %scale0, %scale1 : vector<2xf32> + vector.store %scales, %scale_stage_view[%subgroup, %bounded_block_lane, %c0] : vector<2xf32>, view<4x32x2xf32> + } + %local_sums = vector.from_elements %local_sum0, %local_sum1 : vector<2xf32> + %sums = kernel.subgroup.reduce %local_sums : vector<2xf32> + %sum0 = vector.extract %sums[0] : vector<2xf32> -> f32 + %sum1 = vector.extract %sums[1] : vector<2xf32> -> f32 + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %output_tile = [%c0 to %output_tile_count step %c1] unroll { + %output_tile_channel = index.mul %output_tile, %c128 : index + %lane_output_channel = index.add %output_tile_channel, %lane_output_base : index + %lane_output_channel_valid = index.cmp ult, %lane_output_channel, %value_head_size : index + %safe_lane_output_channel = scf.select %lane_output_channel_valid, %lane_output_channel, %c0 : index + %unnormalized_output0, %unnormalized_output1 = scf.for %block = [%c0 to %active_block_count step %c1](%running_output0 = %c0_f32x2 : vector<2xf32>, %running_output1 = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) unroll(%c4) schedule(interleaved) { + %scales = vector.load %scale_stage_view[%subgroup, %block, %c0] : view<4x32x2xf32> -> vector<2xf32> + %scale0 = vector.extract %scales[0] : vector<2xf32> -> f32 + %scale1 = vector.extract %scales[1] : vector<2xf32> -> f32 + %block_output0_f16 = vector.load %partial_output_view[%key_value_head, %block, %safe_query_row0, %safe_lane_output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> -> vector<2xf16> + %block_output0 = vector.extf %block_output0_f16 : vector<2xf16> to vector<2xf32> + %scale_vector0 = vector.splat %scale0 : vector<2xf32> + %scaled_output0 = vector.mulf %block_output0, %scale_vector0 : vector<2xf32> + %next_output0 = vector.addf %running_output0, %scaled_output0 : vector<2xf32> + %block_output1_f16 = vector.load %partial_output_view[%key_value_head, %block, %safe_query_row1, %safe_lane_output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x[%value_head_size]xf16> -> vector<2xf16> + %block_output1 = vector.extf %block_output1_f16 : vector<2xf16> to vector<2xf32> + %scale_vector1 = vector.splat %scale1 : vector<2xf32> + %scaled_output1 = vector.mulf %block_output1, %scale_vector1 : vector<2xf32> + %next_output1 = vector.addf %running_output1, %scaled_output1 : vector<2xf32> + scf.yield %next_output0, %next_output1 : vector<2xf32>, vector<2xf32> + } + template.apply<@ggml.flash_attention.decode_split.publish_f32.vector2>(%query_head_count, %value_head_size, %query_head_valid0, %query_head_valid1, %lane_output_channel_valid, %query_head0, %query_head1, %lane_output_channel, %sum0, %sum1, %unnormalized_output0, %unnormalized_output1, %output_aligned) : (index, index, i1, i1, i1, index, index, index, f32, f32, vector<2xf32>, vector<2xf32>, buffer) + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/online_softmax.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/online_softmax.loom new file mode 100644 index 000000000000..a6ee1946480c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/online_softmax.loom @@ -0,0 +1,29 @@ +template.decl @ggml.flash_attention.online_softmax.prefill(%masked_scores: vector<4xf32>, %valid: vector<4xi1>, %current_max: vector<4xf32>, %current_sum: vector<4xf32>, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %lane: index, %probability_stage: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) + +template.def<@ggml.flash_attention.online_softmax.prefill> device @ggml_flash_attention_online_softmax_prefill(%masked_scores: vector<4xf32>, %valid: vector<4xi1>, %current_max: vector<4xf32>, %current_sum: vector<4xf32>, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %lane: index, %probability_stage: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %probability_layout = encoding.layout.strided [72, 1] : encoding + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_layout> + %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> + %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> + %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> + %raw_probability = vector.expf %score_delta : vector<4xf32> + %probability = vector.select %valid, %raw_probability, %c0_f32x4 : vector<4xf32> + %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> + %old_delta = vector.subf %current_max, %next_max : vector<4xf32> + %old_scale = vector.expf %old_delta : vector<4xf32> + %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> + %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> + %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> + %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 + %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 + %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 + %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 + view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_layout> + view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_layout> + view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_layout> + view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_layout> + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return %next_max, %next_sum, %old_scale : vector<4xf32>, vector<4xf32>, vector<4xf32> +} 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 new file mode 100644 index 000000000000..71dd2cc24799 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_f32.loom @@ -0,0 +1,176 @@ +config.decl @ggml.flash_attention.apply_gate : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.flash_attention.gate_stride_head : %value: index where [range(%value, 1, 16777216)] + +config.decl @ggml.flash_attention.gate_stride_token : %value: index where [range(%value, 1, 16777216)] + +template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) + +template.decl @ggml.flash_attention.decode_split.publish_f32.scalar(%query_head_count: index, %value_head_size: index, %query_head: index, %output_channel: index, %sum: f32, %unnormalized_output: f32, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.publish_f32.vector2(%query_head_count: index, %value_head_size: index, %query_head_valid0: i1, %query_head_valid1: i1, %lane_output_channel_valid: i1, %query_head0: index, %query_head1: index, %lane_output_channel: index, %sum0: f32, %sum1: f32, %unnormalized_output0: vector<2xf32>, %unnormalized_output1: vector<2xf32>, %output: buffer) + +template.def<@ggml.flash_attention.prefill.publish_output> device @ggml_flash_attention_prefill_publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c1_f32 = scalar.constant 1.0 : f32 + %c1_f32x4 = vector.constant 1.0 : vector<4xf32> + %apply_gate_index = config.get @ggml.flash_attention.apply_gate : index + %gate_stride_head = config.get @ggml.flash_attention.gate_stride_head : index + %gate_stride_token = config.get @ggml.flash_attention.gate_stride_token : index + %apply_gate = index.cmp eq, %apply_gate_index, %c1 : index + %gate_aligned = buffer.assume.alignment %gate {minimum_alignment = 16} : buffer + %gate_view = buffer.view %gate_aligned[%c0_offset] : buffer -> view<268435456xf32> + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf32> + %f16_output_view = buffer.view %f16_output[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + // 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 + %sum3_scalar = vector.extract %final_sum[3] : vector<4xf32> -> f32 + %inverse_sum0_f32 = scalar.divf %c1_f32, %sum0_scalar : f32 + %inverse_sum1_f32 = scalar.divf %c1_f32, %sum1_scalar : f32 + %inverse_sum2_f32 = scalar.divf %c1_f32, %sum2_scalar : f32 + %inverse_sum3_f32 = scalar.divf %c1_f32, %sum3_scalar : f32 + %inverse_sum0_f16 = scalar.fptrunc %inverse_sum0_f32 : f32 to f16 + %inverse_sum1_f16 = scalar.fptrunc %inverse_sum1_f32 : f32 to f16 + %inverse_sum2_f16 = scalar.fptrunc %inverse_sum2_f32 : f32 to f16 + %inverse_sum3_f16 = scalar.fptrunc %inverse_sum3_f32 : f32 to f16 + %inverse_sum0 = vector.splat %inverse_sum0_f16 : vector<4xf16> + %inverse_sum1 = vector.splat %inverse_sum1_f16 : vector<4xf16> + %inverse_sum2 = vector.splat %inverse_sum2_f16 : vector<4xf16> + %inverse_sum3 = vector.splat %inverse_sum3_f16 : vector<4xf16> + %normalized0_f16 = vector.mulf %final_output0, %inverse_sum0 : vector<4xf16> + %normalized1_f16 = vector.mulf %final_output1, %inverse_sum1 : vector<4xf16> + %normalized2_f16 = vector.mulf %final_output2, %inverse_sum2 : vector<4xf16> + %normalized3_f16 = vector.mulf %final_output3, %inverse_sum3 : vector<4xf16> + %normalized0 = vector.extf %normalized0_f16 : vector<4xf16> to vector<4xf32> + %normalized1 = vector.extf %normalized1_f16 : vector<4xf16> to vector<4xf32> + %normalized2 = vector.extf %normalized2_f16 : vector<4xf16> to vector<4xf32> + %normalized3 = vector.extf %normalized3_f16 : vector<4xf16> to vector<4xf32> + scf.if %query_valid0 { + %published0 = scf.if %apply_gate -> (vector<4xf32>) { + %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, %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> + %gate_exp0 = vector.expf %negative_gate0 : vector<4xf32> + %gate_denominator0 = vector.addf %c1_f32x4, %gate_exp0 : vector<4xf32> + %sigmoid_gate0 = vector.divf %c1_f32x4, %gate_denominator0 : vector<4xf32> + %gated_output0 = vector.mulf %normalized0, %sigmoid_gate0 : vector<4xf32> + scf.yield %gated_output0 : vector<4xf32> + } else { + scf.yield %normalized0 : vector<4xf32> + } + 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, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + } + } + scf.if %query_valid1 { + %published1 = scf.if %apply_gate -> (vector<4xf32>) { + %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, %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> + %gate_exp1 = vector.expf %negative_gate1 : vector<4xf32> + %gate_denominator1 = vector.addf %c1_f32x4, %gate_exp1 : vector<4xf32> + %sigmoid_gate1 = vector.divf %c1_f32x4, %gate_denominator1 : vector<4xf32> + %gated_output1 = vector.mulf %normalized1, %sigmoid_gate1 : vector<4xf32> + scf.yield %gated_output1 : vector<4xf32> + } else { + scf.yield %normalized1 : vector<4xf32> + } + 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, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + } + } + scf.if %query_valid2 { + %published2 = scf.if %apply_gate -> (vector<4xf32>) { + %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, %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> + %gate_exp2 = vector.expf %negative_gate2 : vector<4xf32> + %gate_denominator2 = vector.addf %c1_f32x4, %gate_exp2 : vector<4xf32> + %sigmoid_gate2 = vector.divf %c1_f32x4, %gate_denominator2 : vector<4xf32> + %gated_output2 = vector.mulf %normalized2, %sigmoid_gate2 : vector<4xf32> + scf.yield %gated_output2 : vector<4xf32> + } else { + scf.yield %normalized2 : vector<4xf32> + } + 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, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + } + } + scf.if %query_valid3 { + %published3 = scf.if %apply_gate -> (vector<4xf32>) { + %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, %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> + %gate_exp3 = vector.expf %negative_gate3 : vector<4xf32> + %gate_denominator3 = vector.addf %c1_f32x4, %gate_exp3 : vector<4xf32> + %sigmoid_gate3 = vector.divf %c1_f32x4, %gate_denominator3 : vector<4xf32> + %gated_output3 = vector.mulf %normalized3, %sigmoid_gate3 : vector<4xf32> + scf.yield %gated_output3 : vector<4xf32> + } else { + scf.yield %normalized3 : vector<4xf32> + } + 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, %safe_lane_output_channel] : vector<4xf16>, view<[%query_token_count]x[%query_head_count]x[%value_head_size]xf16> + } + } + } + template.return +} + +template.def<@ggml.flash_attention.decode_split.publish_f32.scalar> device @ggml_flash_attention_decode_split_publish_f32_scalar(%query_head_count: index, %value_head_size: index, %query_head: index, %output_channel: index, %sum: f32, %unnormalized_output: f32, %output: buffer) { + %c0_offset = index.constant 0 : offset + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%query_head_count]x[%value_head_size]xf32> + %normalized_output = scalar.divf %unnormalized_output, %sum : f32 + view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x[%value_head_size]xf32> + template.return +} + +template.def<@ggml.flash_attention.decode_split.publish_f32.vector2> device @ggml_flash_attention_decode_split_publish_f32_vector2(%query_head_count: index, %value_head_size: index, %query_head_valid0: i1, %query_head_valid1: i1, %lane_output_channel_valid: i1, %query_head0: index, %query_head1: index, %lane_output_channel: index, %sum0: f32, %sum1: f32, %unnormalized_output0: vector<2xf32>, %unnormalized_output1: vector<2xf32>, %output: buffer) { + %c0_offset = index.constant 0 : offset + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%query_head_count]x[%value_head_size]xf32> + %sum_vector0 = vector.splat %sum0 : vector<2xf32> + %sum_vector1 = vector.splat %sum1 : vector<2xf32> + %normalized_output0 = vector.divf %unnormalized_output0, %sum_vector0 : vector<2xf32> + %normalized_output1 = vector.divf %unnormalized_output1, %sum_vector1 : vector<2xf32> + %publish_output0 = scalar.andi %query_head_valid0, %lane_output_channel_valid : i1 + %publish_output1 = scalar.andi %query_head_valid1, %lane_output_channel_valid : i1 + scf.if %publish_output0 { + vector.store %normalized_output0, %output_view[%query_head0, %lane_output_channel] : vector<2xf32>, view<[%query_head_count]x[%value_head_size]xf32> + } + scf.if %publish_output1 { + vector.store %normalized_output1, %output_view[%query_head1, %lane_output_channel] : vector<2xf32>, view<[%query_head_count]x[%value_head_size]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_q8.loom new file mode 100644 index 000000000000..7505c89d2dc9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/publish_q8.loom @@ -0,0 +1,83 @@ +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] + +template.decl @ggml.flash_attention.decode_split.pack_completed_q8(%key_value_head: index, %output: buffer, %q8_output: buffer) + +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) + +// Packs the query heads owned by one completed KV head into GGML's Q8_1 x4 +// layout. Each 32-workitem cohort owns one contiguous 128-element query row; +// the production 8:1 GQA ratio therefore fills all 256 workitems. Smaller or +// larger valid ratios use the same phase loop without making barriers +// conditional. +template.def<@ggml.flash_attention.decode_split.pack_completed_q8> device @ggml_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head: index, %output: buffer, %q8_output: buffer) { + %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 + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %output_noalias, %q8_output_noalias = buffer.assume.noalias %output, %q8_output : buffer, buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %q8_output_aligned = buffer.assume.alignment %q8_output_noalias {minimum_alignment = 16} : buffer + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x[%value_head_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> + %padded_value_head_size = index.add %value_head_size, %c127 : index + %head_tile_count = index.div %padded_value_head_size, %c128 : index + %group_in_phase = index.div %workitem, %c32 : index + %lane_in_group0 = index.rem %workitem, %c32 : index + %lane_in_group = index.assume %lane_in_group0 [range(%lane_in_group0, 0, 31)] : index + %block_in_group0 = index.div %lane_in_group, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane_in_group, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %padded_phase_count = index.add %query_heads_per_key_value_head, %c7 : index + %phase_count = index.div %padded_phase_count, %c8 : index + scf.for %phase = [%c0 to %phase_count step %c1] unroll { + %phase_group_base = index.mul %phase, %c8 : index + %group_in_key_value_head = index.add %phase_group_base, %group_in_phase : index + %valid_group = index.cmp ult, %group_in_key_value_head, %query_heads_per_key_value_head : index + %key_value_query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index + %query_head0 = index.add %key_value_query_head_base, %group_in_key_value_head : index + %safe_query_head = scf.select %valid_group, %query_head0, %c0 : index + scf.for %head_tile = [%c0 to %head_tile_count step %c1] unroll { + %head_tile_channel = index.mul %head_tile, %c128 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_channel_add = index.add %block_element_add, %word_element_add : index + %input_channel0 = index.add %head_tile_channel, %input_channel_add : index + %input_channel = index.assume %input_channel0 [range(%input_channel0, 0, 508)] : index + %valid_channel = index.cmp ult, %input_channel, %value_head_size : index + %publish_word = scalar.andi %valid_group, %valid_channel : i1 + %input_values = scf.if %publish_word -> (vector<4xf32>) { + %values = vector.load %output_view[%safe_query_head, %input_channel] : view<[%query_head_count]x[%value_head_size]xf32> -> vector<4xf32> + scf.yield %values : vector<4xf32> + } else { + scf.yield %c0_f32x4 : vector<4xf32> + } + %flattened_head_channel = index.mul %safe_query_head, %value_head_size : index + %flattened_channel = index.add %flattened_head_channel, %input_channel : index + template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%publish_word, %c0_offset, %flattened_channel, %input_values, %scratch_values, %scratch_d, %q8_output_aligned) : (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/flash_attention/pv_block_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/pv_block_wmma.loom new file mode 100644 index 000000000000..9008a1dce89a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/pv_block_wmma.loom @@ -0,0 +1,143 @@ +template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + +template.decl @ggml.flash_attention.pv_block_wmma.prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + +template.def<@ggml.flash_attention.pv_block_wmma.prefill_full> device @ggml_flash_attention_pv_block_wmma_prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %probability_layout = encoding.layout.strided [72, 1] : encoding + %value_layout = encoding.layout.strided [%value_row_stride, %value_column_stride] : encoding + %value_view = buffer.view %value[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c4 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { + %output_tile_valid = index.cmp ult, %output_tile, %group_output_tile_count : index + %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %output_tile_valid -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %global_output_tile0 = index.add %output_group_tile_base, %output_tile : index + %global_output_tile = index.assume %global_output_tile0 [lt(%global_output_tile0, %output_tile_count)] : index + %output_tile_channel = index.mul %global_output_tile, %c64 : index + %value_channel0 = index.add %value_head_base, %output_tile_channel : index + %value_channel = index.add %value_channel0, %subgroup_product_channel : index + %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll schedule(recurrence) { + %value_token0 = index.add %key_origin, %key_tile : index + %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index + %tile_updated_output0, %tile_updated_output1, %tile_updated_output2, %tile_updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> + %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> + %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> + %updated_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> + scf.yield %updated_tile_output0, %updated_tile_output1, %updated_tile_output2, %updated_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // Complete every read before the next output tile overwrites LDS. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %tile_updated_output0, %tile_updated_output1, %tile_updated_output2, %tile_updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + template.return %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> +} + +template.def<@ggml.flash_attention.pv_block_wmma.prefill_tail> device @ggml_flash_attention_pv_block_wmma_prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %probability_layout = encoding.layout.strided [72, 1] : encoding + %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x[%tail_stage_head_size]xf16> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_layout> + %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> + %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> + %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 + %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 + %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 + %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 + %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> + %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> + %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> + %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> + %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> + %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> + %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> + %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> + %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c4 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { + %output_tile_valid = index.cmp ult, %output_tile, %group_output_tile_count : index + %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %output_tile_valid -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %output_tile_channel = index.mul %output_tile, %c64 : index + %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index + %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> + %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { + %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_layout> -> vector<16xf16> + %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x[%tail_stage_head_size]xf16> -> vector<16xf16> + %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %next_product_accumulator : vector<8xf16> + } + vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index + %tile_updated_output0, %tile_updated_output1, %tile_updated_output2, %tile_updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> + %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> + %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> + %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> + %updated_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> + scf.yield %updated_tile_output0, %updated_tile_output1, %updated_tile_output2, %updated_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %tile_updated_output0, %tile_updated_output1, %tile_updated_output2, %tile_updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + template.return %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom new file mode 100644 index 000000000000..24449ad70d5e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/flash_attention/qk_block_wmma.loom @@ -0,0 +1,436 @@ +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.qk_head_size : %value: index where [range(%value, 16, 512), mul(%value, 16)] + +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] + +config.decl @ggml.flash_attention.attention_scale : f32 + +config.def @ggml.flash_attention.value_layout = 0 : index + +template.decl @ggml.flash_attention.prefill.stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer) + +template.decl @ggml.flash_attention.qk_block_wmma.prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer) + +template.decl @ggml.flash_attention.qk_block_wmma.prefill_tail(%qk_head_size: index, %query_stage_stride: index, %tail_stage_head_size: index, %tail_score_wave: i1, %subgroup: index, %tail_key_value_stage: buffer, %query_stage: buffer, %score_stage: buffer) + +template.decl @ggml.flash_attention.online_softmax.prefill(%masked_scores: vector<4xf32>, %valid: vector<4xi1>, %current_max: vector<4xf32>, %current_sum: vector<4xf32>, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %lane: index, %probability_stage: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) + +template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + +template.decl @ggml.flash_attention.pv_block_wmma.prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + +template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) + +template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) + +template.def<@ggml.flash_attention.prefill.stage_query> device @ggml_flash_attention_prefill_stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer) { + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f16 = scalar.constant 0.0 : f16 + %query_view = buffer.view %query[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x[%qk_head_size]xf32> + %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x[%query_stage_stride]xf16> + // Scale and truncate Q exactly once. The Vulkan reference does this before + // entering its KV loop, making QK a native F16 WMMA while retaining F32 + // accumulation. + scf.for %load_iteration = [%c0 to %query_load_iteration_count step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %local_query_row = index.div %linear, %qk_head_size : index + %query_channel = index.rem %linear, %qk_head_size : index + %query_token = index.add %query_origin, %local_query_row : index + %query_valid_load = index.cmp ult, %query_token, %query_token_count : index + %query_value = scf.if %query_valid_load -> (f16) { + %loaded = view.load %query_view[%query_token, %query_head, %query_channel] : view<[%query_token_count]x[%query_head_count]x[%qk_head_size]xf32> -> f32 + %scaled = scalar.mulf %loaded, %attention_scale : f32 + %truncated = scalar.fptrunc %scaled : f32 to f16 + scf.yield %truncated : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x[%query_stage_stride]xf16> + } + template.return +} + +template.def<@ggml.flash_attention.qk_block_wmma.prefill_full> device @ggml_flash_attention_qk_block_wmma_prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer) { + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %c0_offset = index.constant 0 : offset + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %query_transposed_layout = encoding.layout.strided [1, %query_stage_stride] : encoding + %key_view = buffer.view %key[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x20xf32> + %score_init_values = vector.constant 0.0 : vector<4xf32> + %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> + %score_fragment = scf.for %head_tile = [%c0 to %qk_head_size step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll schedule(recurrence) { + %key_channel = index.add %key_head_base, %head_tile : index + %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%bounded_key_value_token_count]x[%key_width]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<[%qk_head_size]x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x20xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return +} + +template.def<@ggml.flash_attention.prefill.body> device @ggml_flash_attention_prefill_body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) { + %bounded_query_token_count = index.assume %query_token_count [range(%query_token_count, 1, 2048)] : index + %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 + %qk_head_size0 = config.get @ggml.flash_attention.qk_head_size : index + %qk_head_size = index.assume %qk_head_size0 [range(%qk_head_size0, 16, 512), mul(%qk_head_size0, 16)] : index + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %query_tile = kernel.workgroup.id : index + %query_head0 = kernel.workgroup.id : index + %output_group0 = kernel.workgroup.id : index + %query_head, %launch_query_head_count = index.assume %query_head0, %query_head_count [lt(%query_head0, %query_head_count)] : index, index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %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 + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %f16_bytes = index.constant 2 : offset + %score_stage_bytes = index.constant 5120 : offset + %probability_stage_bytes = index.constant 2304 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %attention_scale = config.get @ggml.flash_attention.attention_scale : f32 + %c0_f16 = scalar.constant 0.0 : f16 + %output_zero0 = vector.constant 0.0 : vector<4xf16> + %output_zero1 = vector.constant 0.0 : vector<4xf16> + %output_zero2 = vector.constant 0.0 : vector<4xf16> + %output_zero3 = vector.constant 0.0 : vector<4xf16> + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %query_heads_per_key_value_head = index.div %query_head_count, %key_value_head_count : index + %key_value_head = index.div %query_head, %query_heads_per_key_value_head : index + %key_width = index.mul %key_value_head_count, %qk_head_size : index + %value_width = index.mul %key_value_head_count, %value_head_size : index + %key_head_base = index.mul %key_value_head, %qk_head_size : index + %value_head_base = index.mul %key_value_head, %value_head_size : index + %output_tile_count = index.div %value_head_size, %c64 : index + %padded_output_tile_count = index.add %output_tile_count, %c3 : index + %output_group_count = index.div %padded_output_tile_count, %c4 : index + %output_group, %launch_output_group_count = index.assume %output_group0, %output_group_count [lt(%output_group0, %output_group_count)] : index, index + %output_group_tile_base = index.mul %output_group, %c4 : index + %remaining_output_tile_count = index.sub %output_tile_count, %output_group_tile_base : index + %group_output_tile_count = index.min %remaining_output_tile_count, %c4 : index + %group_output_channel_count = index.mul %group_output_tile_count, %c64 : index + %query_stage_stride = index.add %qk_head_size, %c8 : index + %query_stage_element_count = index.mul %c16, %query_stage_stride : index + %query_stage_bytes = index.scale %query_stage_element_count, %f16_bytes : index, offset -> offset + %query_element_count = index.mul %c16, %qk_head_size : index + %query_load_iteration_count = index.div %query_element_count, %c256 : index + %tail_stage_head_size_is_value = index.cmp ult, %qk_head_size, %value_head_size : index + %tail_stage_head_size = scf.select %tail_stage_head_size_is_value, %value_head_size, %qk_head_size : index + %tail_key_element_count = index.mul %c32, %qk_head_size : index + %tail_key_load_iteration_count = index.div %tail_key_element_count, %c256 : index + %tail_value_stage_channel_count = index.min %value_head_size, %c256 : index + %tail_value_element_count = index.mul %c32, %tail_value_stage_channel_count : index + %tail_value_load_iteration_count = index.div %tail_value_element_count, %c256 : index + %tail_key_value_stage_element_count = index.mul %c32, %tail_stage_head_size : index + %tail_key_value_stage_capacity = index.scale %tail_key_value_stage_element_count, %f16_bytes : index, offset -> offset + %query_origin = index.mul %query_tile, %c16 : index + %full_key_value_block_count = index.div %bounded_key_value_token_count, %c64 : index + %full_key_value_token_count0 = index.mul %full_key_value_block_count, %c64 : index + %full_key_value_token_count = index.assume %full_key_value_token_count0 [range(%full_key_value_token_count0, 0, 32768), mul(%full_key_value_token_count0, 64)] : index + %tail_key_value_token_count = index.sub %bounded_key_value_token_count, %full_key_value_token_count : index + %has_key_value_tail = index.cmp ne, %tail_key_value_token_count, %c0 : index + %tail_key_value_stage_bytes = scf.select %has_key_value_tail, %tail_key_value_stage_capacity, %c0_offset : offset + %subgroup_score_column = index.mul %subgroup, %c16 : index + %subgroup_query_row = index.mul %subgroup, %c4 : index + %query_row0 = index.add %subgroup_query_row, %c0 : index + %query_row1 = index.add %subgroup_query_row, %c1 : index + %query_row2 = index.add %subgroup_query_row, %c2 : index + %query_row3 = index.add %subgroup_query_row, %c3 : index + %query_token0 = index.add %query_origin, %query_row0 : index + %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, %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 + %lane_product_channel0 = index.rem %lane, %c16 : index + %lane_product_channel = index.mul %lane_product_channel0, %c4 : index + %output_group_channel_base = index.mul %output_group, %c256 : index + %lane_output_channel0 = index.mul %lane, %c4 : index + %lane_output_channel = index.add %output_group_channel_base, %lane_output_channel0 : index + %lane_output_end = index.add %lane_output_channel, %c4 : index + %lane_has_output = index.cmp ule, %lane_output_end, %value_head_size : index + // Q and score rows retain padding for their ownership exchanges. Store + // probabilities query-major with eight spare K columns so native LHS + // fragments use aligned wide reads instead of scalar transposed loads. + %query_transposed_layout = encoding.layout.strided [1, %query_stage_stride] : encoding + %probability_layout = encoding.layout.strided [72, 1] : encoding + %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer + %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer + %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer + %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer + %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_width]xf16> + %value_layout_tag = config.get @ggml.flash_attention.value_layout : index + %transposed_value = index.cmp eq, %value_layout_tag, %c1 : index + %value_row_stride = scf.select %transposed_value, %c1, %value_width : index + %value_column_stride = scf.select %transposed_value, %bounded_key_value_token_count, %c1 : index + %value_layout = encoding.layout.strided [%value_row_stride, %value_column_stride] : encoding + %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> + %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> + %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 + %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x20xf32> + %mask_summary_view = buffer.view %score_stage[%c0_offset] : buffer -> view<4xf32> + %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_layout> + // Score reads retire at probability publication. Product reads retire + // before the next mask/score phase, so the two exchanges share storage. + %product_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<16x64xf16> + %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x[%tail_stage_head_size]xf16> + template.apply<@ggml.flash_attention.prefill.stage_query>(%bounded_query_token_count, %query_origin, %query_head, %query_load_iteration_count, %query_stage_stride, %qk_head_size, %attention_scale, %query_aligned, %query_stage) : (index, index, index, index, index, index, f32, buffer, buffer) + 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 + // KV blocks entirely at negative infinity; reducing the cached tile lets + // the workgroup skip QK, softmax, and P*V for those blocks. This mirrors + // the Vulkan oracle without requiring its auxiliary compact mask buffer. + %key_token0 = index.add %key_origin, %lane : index + %key_token = index.assume %key_token0 [lt(%key_token0, %bounded_key_value_token_count)] : index + %mask0 = scf.if %query_valid0 -> (f16) { + %value = view.load %mask_view[%query_token0, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask1 = scf.if %query_valid1 -> (f16) { + %value = view.load %mask_view[%query_token1, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask2 = scf.if %query_valid2 -> (f16) { + %value = view.load %mask_view[%query_token2, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask3 = scf.if %query_valid3 -> (f16) { + %value = view.load %mask_view[%query_token3, %key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + scf.yield %value : f16 + } else { + scf.yield %c0_f16 : f16 + } + %mask_f16 = vector.from_elements %mask0, %mask1, %mask2, %mask3 : vector<4xf16> + %mask_summary_f32 = vector.extf %mask_f16 : vector<4xf16> to vector<4xf32> + %effective_mask = vector.select %query_valid, %mask_summary_f32, %negative_f32x4 : vector<4xf32> + %lane_mask_maximum = vector.reduce %effective_mask, %negative_large : vector<4xf32>, f32 + %subgroup_mask_maximum = kernel.subgroup.reduce %lane_mask_maximum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_mask_maximum, %mask_summary_view[%subgroup] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %subgroup_mask_maxima = vector.load %mask_summary_view[%c0] : view<4xf32> -> vector<4xf32> + %workgroup_mask_maximum = vector.reduce %subgroup_mask_maxima, %negative_large : vector<4xf32>, f32 + %block_has_attention = scalar.cmpf ogt, %workgroup_mask_maximum, %negative_large : f32 + %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 = scf.if %block_has_attention -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + // Four independent wave-level WMMAs produce a 16x64 score tile. + %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index + %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index + %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index + template.apply<@ggml.flash_attention.qk_block_wmma.prefill_full>(%bounded_key_value_token_count, %key_width, %qk_head_size, %query_stage_stride, %score_key_origin, %key_head_base, %subgroup_score_column, %key_aligned, %query_stage, %score_stage) : (index, index, index, index, index, index, index, buffer, buffer, buffer) + // LDS transposes ownership from one 16-column score slice per wave to + // four complete query rows per wave. Every lane contributes one key + // column to each of those rows. + %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x20xf32> -> f32 + %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x20xf32> -> f32 + %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x20xf32> -> f32 + %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x20xf32> -> f32 + %mask_f32 = vector.extf %mask_f16 : vector<4xf16> to vector<4xf32> + %raw_scores = vector.from_elements %raw_score0, %raw_score1, %raw_score2, %raw_score3 : vector<4xf32> + %masked_scores0 = vector.addf %raw_scores, %mask_f32 : vector<4xf32> + %masked_scores = vector.select %query_valid, %masked_scores0, %negative_f32x4 : vector<4xf32> + %next_max, %next_sum, %old_scale = template.apply<@ggml.flash_attention.online_softmax.prefill>(%masked_scores, %query_valid, %current_max, %current_sum, %query_row0, %query_row1, %query_row2, %query_row3, %lane, %probability_stage) : (vector<4xf32>, vector<4xi1>, vector<4xf32>, vector<4xf32>, index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) + // Compute the 64-channel tiles in this output group sequentially. This + // keeps one P*V accumulator live per wave and reuses a 2 KiB exchange + // tile instead of retaining multiple output groups. + %next_output0, %next_output1, %next_output2, %next_output3 = template.apply<@ggml.flash_attention.pv_block_wmma.prefill_full>(%old_scale, %current_output0, %current_output1, %current_output2, %current_output3, %group_output_tile_count, %output_group_tile_base, %output_tile_count, %value_head_base, %subgroup_product_channel, %lane_output_tile, %lane_product_channel, %query_row0, %query_row1, %query_row2, %query_row3, %key_origin, %last_full_key_tile_start, %bounded_key_value_token_count, %value_width, %value_row_stride, %value_column_stride, %value_aligned, %probability_stage, %score_stage) : (vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } else { + scf.yield %current_max, %current_sum, %current_output0, %current_output1, %current_output2, %current_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + scf.yield %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + // A masked 32-row tile handles the final 1-63 KV rows with the same WMMA + // fragments as the aligned path. K and V are staged separately into product + // scratch, so every padded element is initialized and no physical padding is + // required of the caller. Each tile rounds probabilities to F16 before P*V. + %tail_score_wave = index.cmp ult, %subgroup, %c2 : index + %final_max, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3 = scf.for %tail_key_origin = [%full_key_value_token_count to %bounded_key_value_token_count step %c32](%current_max = %full_max : vector<4xf32>, %current_sum = %full_sum : vector<4xf32>, %current_output0 = %full_output0 : vector<4xf16>, %current_output1 = %full_output1 : vector<4xf16>, %current_output2 = %full_output2 : vector<4xf16>, %current_output3 = %full_output3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { + %tail_remaining = index.sub %bounded_key_value_token_count, %tail_key_origin : index + %tail_key_count = index.min %tail_remaining, %c32 : index + // Cooperatively stage one K tile, explicitly zeroing the padded rows. + scf.for %load_iteration = [%c0 to %tail_key_load_iteration_count step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %tail_key_row = index.div %linear, %qk_head_size : index + %tail_key_channel = index.rem %linear, %qk_head_size : index + %tail_key_valid = index.cmp ult, %tail_key_row, %tail_key_count : index + %tail_key_value = scf.if %tail_key_valid -> (f16) { + %tail_key_token0 = index.add %tail_key_origin, %tail_key_row : index + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %global_key_channel = index.add %key_head_base, %tail_key_channel : index + %loaded = view.load %key_view[%tail_key_token, %global_key_channel] : view<[%bounded_key_value_token_count]x[%key_width]xf16> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %tail_key_value, %tail_key_value_stage_view[%tail_key_row, %tail_key_channel] : f16, view<32x[%tail_stage_head_size]xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Waves zero and one compute the two 16-column QK fragments in this tile. + template.apply<@ggml.flash_attention.qk_block_wmma.prefill_tail>(%qk_head_size, %query_stage_stride, %tail_stage_head_size, %tail_score_wave, %subgroup, %tail_key_value_stage, %query_stage, %score_stage) : (index, index, index, i1, index, buffer, buffer, buffer) + // LDS changes ownership from the QK wave to four query-row waves. Lanes + // beyond the logical tail never read the score or mask buffers. + %tail_lane_valid = index.cmp ult, %lane, %tail_key_count : index + %tail_valid0 = scalar.andi %tail_lane_valid, %query_valid0 : i1 + %tail_valid1 = scalar.andi %tail_lane_valid, %query_valid1 : i1 + %tail_valid2 = scalar.andi %tail_lane_valid, %query_valid2 : i1 + %tail_valid3 = scalar.andi %tail_lane_valid, %query_valid3 : i1 + %tail_valid = vector.from_elements %tail_valid0, %tail_valid1, %tail_valid2, %tail_valid3 : vector<4xi1> + %tail_key_token0 = index.add %tail_key_origin, %lane : index + %masked_score0 = scf.if %tail_valid0 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row0] : view<64x20xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token0, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score1 = scf.if %tail_valid1 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row1] : view<64x20xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token1, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score2 = scf.if %tail_valid2 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row2] : view<64x20xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token2, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_score3 = scf.if %tail_valid3 -> (f32) { + %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index + %raw_score = view.load %score_stage_view[%lane, %query_row3] : view<64x20xf32> -> f32 + %mask_f16 = view.load %mask_view[%query_token3, %tail_key_token] : view<[%bounded_query_token_count]x[%bounded_key_value_token_count]xf16> -> f16 + %mask_f32 = scalar.extf %mask_f16 : f16 to f32 + %score = scalar.addf %raw_score, %mask_f32 : f32 + scf.yield %score : f32 + } else { + scf.yield %negative_large : f32 + } + %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> + %next_max, %next_sum, %old_scale = template.apply<@ggml.flash_attention.online_softmax.prefill>(%masked_scores, %tail_valid, %current_max, %current_sum, %query_row0, %query_row1, %query_row2, %query_row3, %lane, %probability_stage) : (vector<4xf32>, vector<4xi1>, vector<4xf32>, vector<4xf32>, index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) + // Reuse the tail K scratch for V after probability publication. Each + // workgroup stages only its output group, starting at local column zero. + // Stage the fixed maximum output group and zero-pad a partial group. + // This keeps the unrolled trip count and local coordinates static. + scf.for %load_iteration = [%c0 to %tail_value_load_iteration_count step %c1] unroll { + %linear = index.madd %load_iteration, %c256, %workitem : index + %tail_value_row = index.div %linear, %tail_value_stage_channel_count : index + %tail_value_channel = index.rem %linear, %tail_value_stage_channel_count : index + %tail_value_row_valid = index.cmp ult, %tail_value_row, %tail_key_count : index + %tail_value_channel_valid = index.cmp ult, %tail_value_channel, %group_output_channel_count : index + %tail_value_valid = scalar.andi %tail_value_row_valid, %tail_value_channel_valid : i1 + %output_value_channel = index.add %output_group_channel_base, %tail_value_channel : index + %tail_value = scf.if %tail_value_valid -> (f16) { + %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index + %tail_value_token = index.assume %tail_value_token0 [lt(%tail_value_token0, %bounded_key_value_token_count)] : index + %global_value_channel0 = index.add %value_head_base, %output_value_channel : index + %global_value_channel = index.assume %global_value_channel0 [lt(%global_value_channel0, %value_width)] : index + %loaded = view.load %value_view[%tail_value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%value_width]xf16, %value_layout> -> f16 + scf.yield %loaded : f16 + } else { + scf.yield %c0_f16 : f16 + } + view.store %tail_value, %tail_key_value_stage_view[%tail_value_row, %tail_value_channel] : f16, view<32x[%tail_stage_head_size]xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // Keep tail V staging separate from the 2 KiB product exchange, then use + // the same grouped 64-channel phases as the aligned path. + %next_output0, %next_output1, %next_output2, %next_output3 = template.apply<@ggml.flash_attention.pv_block_wmma.prefill_tail>(%old_scale, %current_output0, %current_output1, %current_output2, %current_output3, %group_output_tile_count, %subgroup_product_channel, %lane_output_tile, %lane_product_channel, %query_row0, %query_row1, %query_row2, %query_row3, %tail_stage_head_size, %tail_key_value_stage, %probability_stage, %score_stage) : (vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> + } + template.apply<@ggml.flash_attention.prefill.publish_output>(%publish_f16, %bounded_query_token_count, %query_head_count, %value_head_size, %lane_has_output, %query_valid0, %query_valid1, %query_valid2, %query_valid3, %query_token0, %query_token1, %query_token2, %query_token3, %query_head, %lane_output_channel, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3, %gate, %output_aligned, %f16_output) : (i1, index, index, index, i1, i1, i1, i1, i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>, buffer, buffer, buffer) + template.return +} + +template.def<@ggml.flash_attention.qk_block_wmma.prefill_tail> device @ggml_flash_attention_qk_block_wmma_prefill_tail(%qk_head_size: index, %query_stage_stride: index, %tail_stage_head_size: index, %tail_score_wave: i1, %subgroup: index, %tail_key_value_stage: buffer, %query_stage: buffer, %score_stage: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c0_offset = index.constant 0 : offset + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %query_transposed_layout = encoding.layout.strided [1, %query_stage_stride] : encoding + %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x[%tail_stage_head_size]xf16> + %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<[%qk_head_size]x16xf16, %query_transposed_layout> + %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x20xf32> + scf.if %tail_score_wave { + %tail_score_subgroup = index.assume %subgroup [range(%subgroup, 0, 1)] : index + %tail_score_column = index.mul %tail_score_subgroup, %c16 : index + %tail_score_init_values = vector.constant 0.0 : vector<4xf32> + %tail_score_init = vector.fragment %tail_score_init_values shape [%m, %n] : vector<4xf32> + %tail_score_fragment = scf.for %head_tile = [%c0 to %qk_head_size step %c16](%score_accumulator = %tail_score_init : vector<4xf32>) -> (vector<4xf32>) unroll { + %key_fragment = vector.fragment.load %tail_key_value_stage_view[%tail_score_column, %head_tile] shape [%m, %k] : view<32x[%tail_stage_head_size]xf16> -> vector<16xf16> + %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<[%qk_head_size]x16xf16, %query_transposed_layout> -> vector<16xf16> + %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> + scf.yield %next_score_accumulator : vector<4xf32> + } + vector.fragment.store %tail_score_fragment, %score_stage_view[%tail_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x20xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/inplace_state_head128.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/inplace_state_head128.loom new file mode 100644 index 000000000000..7de2c42e34ae --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/inplace_state_head128.loom @@ -0,0 +1,24 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.inplace_state.publish_head128(%state: view<1073741824xf32>, %index0: index, %index1: index, %index2: index, %index3: index, %index4: index, %index5: index, %index6: index, %index7: index, %index8: index, %index9: index, %index10: index, %index11: index, %index12: index, %index13: index, %index14: index, %index15: index, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) + +template.def<@llm.gated_delta_net.inplace_state.publish_head128> device @llm_gated_delta_net_inplace_state_publish_head128(%state: view<1073741824xf32>, %index0: index, %index1: index, %index2: index, %index3: index, %index4: index, %index5: index, %index6: index, %index7: index, %index8: index, %index9: index, %index10: index, %index11: index, %index12: index, %index13: index, %index14: index, %index15: index, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) { + vector.store %value0, %state[%index0] : vector<4xf32>, view<1073741824xf32> + vector.store %value1, %state[%index1] : vector<4xf32>, view<1073741824xf32> + vector.store %value2, %state[%index2] : vector<4xf32>, view<1073741824xf32> + vector.store %value3, %state[%index3] : vector<4xf32>, view<1073741824xf32> + vector.store %value4, %state[%index4] : vector<4xf32>, view<1073741824xf32> + vector.store %value5, %state[%index5] : vector<4xf32>, view<1073741824xf32> + vector.store %value6, %state[%index6] : vector<4xf32>, view<1073741824xf32> + vector.store %value7, %state[%index7] : vector<4xf32>, view<1073741824xf32> + vector.store %value8, %state[%index8] : vector<4xf32>, view<1073741824xf32> + vector.store %value9, %state[%index9] : vector<4xf32>, view<1073741824xf32> + vector.store %value10, %state[%index10] : vector<4xf32>, view<1073741824xf32> + vector.store %value11, %state[%index11] : vector<4xf32>, view<1073741824xf32> + vector.store %value12, %state[%index12] : vector<4xf32>, view<1073741824xf32> + vector.store %value13, %state[%index13] : vector<4xf32>, view<1073741824xf32> + vector.store %value14, %state[%index14] : vector<4xf32>, view<1073741824xf32> + vector.store %value15, %state[%index15] : vector<4xf32>, view<1073741824xf32> + 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 new file mode 100644 index 000000000000..4d6a19003c00 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/projection_epilogue.loom @@ -0,0 +1,69 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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(%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_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) { + %biased = scalar.addf %alpha, %bias : f32 + %abs = scalar.absf %biased : f32 + %neg_abs = scalar.negf %abs : f32 + %exp = scalar.expf %neg_abs : f32 + %one_plus = scalar.addf %one, %exp : f32 + %log = scalar.logf %one_plus : f32 + %positive = scalar.maxnumf %biased, %zero : f32 + %softplus = scalar.addf %positive, %log : f32 + %gate = scalar.mulf %softplus, %a_scale : f32 + %beta_result = scalar.logisticf %beta : f32 + 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(%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 + %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 %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 + %lane0 = kernel.workitem.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 1023)] : index + %linear_base = index.mul %workgroup, %workgroup_size : index + %linear0 = index.add %linear_base, %lane : index + %linear = index.assume %linear0 [range(%linear0, 0, 1048575)] : index + %in_bounds = index.cmp ult, %linear, %elements : index + + %alpha_global = buffer.assume.memory_space %alpha_raw : buffer + %beta_global = buffer.assume.memory_space %beta_raw : buffer + %bias_global = buffer.assume.memory_space %bias : buffer + %a_scale_global = buffer.assume.memory_space %a_scale : buffer + %gate_global = buffer.assume.memory_space %gate_dst : buffer + %beta_dst_global = buffer.assume.memory_space %beta_dst : buffer + %alpha_noalias, %beta_noalias, %bias_noalias, %a_scale_noalias, %gate_noalias, %beta_dst_noalias = buffer.assume.noalias %alpha_global, %beta_global, %bias_global, %a_scale_global, %gate_global, %beta_dst_global : buffer, buffer, buffer, buffer, buffer, buffer + %alpha_view = buffer.view %alpha_noalias[%base] : buffer -> view<1048576xf32> + %beta_view = buffer.view %beta_noalias[%base] : buffer -> view<1048576xf32> + %bias_view = buffer.view %bias_noalias[%base] : buffer -> view<4096xf32> + %a_scale_view = buffer.view %a_scale_noalias[%base] : buffer -> view<4096xf32> + %gate_view = buffer.view %gate_noalias[%base] : buffer -> view<1048576xf32> + %beta_dst_view = buffer.view %beta_dst_noalias[%base] : buffer -> view<1048576xf32> + + scf.if %in_bounds { + %head = index.rem %linear, %heads : index + %alpha = view.load %alpha_view[%linear] : view<1048576xf32> -> f32 + %beta = view.load %beta_view[%linear] : view<1048576xf32> -> f32 + %bias_value = view.load %bias_view[%head] : view<4096xf32> -> f32 + %a_value = view.load %a_scale_view[%head] : view<4096xf32> -> f32 + %c0 = scalar.constant 0.0 : f32 + %c1 = scalar.constant 1.0 : f32 + %gate, %beta_result = template.apply<@llm.gated_delta_net.projection_epilogue.apply>(%alpha, %beta, %bias_value, %a_value, %c0, %c1) : (f32, f32, f32, f32, f32, f32) -> (f32, f32) + template.apply<@llm.gated_delta_net.publish_f32.projection_pair>(%gate, %beta_result, %gate_view, %beta_dst_view, %linear) : (f32, f32, view<1048576xf32>, view<1048576xf32>, index) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/publish_f32.loom new file mode 100644 index 000000000000..83a165dd28ae --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/publish_f32.loom @@ -0,0 +1,73 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.publish_f32.scalar(%value: f32, %output: view<1073741824xf32>, %index: index) + +template.decl @llm.gated_delta_net.publish_f32.projection_pair(%gate: f32, %beta: f32, %gate_output: view<1048576xf32>, %beta_output: view<1048576xf32>, %index: index) + +template.decl @llm.gated_delta_net.publish_f32.state_head128(%column: index, %head_width: index, %state_offset: index, %attention_elements: index, %row: index, %state_stride: index, %output: view<1073741824xf32>, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) + +template.def<@llm.gated_delta_net.publish_f32.scalar> device @llm_gated_delta_net_publish_f32_scalar(%value: f32, %output: view<1073741824xf32>, %index: index) { + view.store %value, %output[%index] : f32, view<1073741824xf32> + template.return +} +template.def<@llm.gated_delta_net.publish_f32.projection_pair> device @llm_gated_delta_net_publish_f32_projection_pair(%gate: f32, %beta: f32, %gate_output: view<1048576xf32>, %beta_output: view<1048576xf32>, %index: index) { + view.store %gate, %gate_output[%index] : f32, view<1048576xf32> + view.store %beta, %beta_output[%index] : f32, view<1048576xf32> + template.return +} + +template.def<@llm.gated_delta_net.publish_f32.state_head128> device @llm_gated_delta_net_publish_f32_state_head128(%column: index, %head_width: index, %state_offset: index, %attention_elements: index, %row: index, %state_stride: index, %output: view<1073741824xf32>, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) { + %woc0_m = index.mul %column, %head_width : index + %woc0_a = index.add %state_offset, %woc0_m : index + %woc0_e = index.add %woc0_a, %attention_elements : index + %woc0_r = index.add %woc0_e, %row : index + %woc0 = index.assume %woc0_r [range(%woc0_r, 0, 1073741820)] : index + vector.store %value0, %output[%woc0] : vector<4xf32>, view<1073741824xf32> + %woc1_r = index.add %woc0, %state_stride : index + %woc1 = index.assume %woc1_r [range(%woc1_r, 0, 1073741820)] : index + vector.store %value1, %output[%woc1] : vector<4xf32>, view<1073741824xf32> + %woc2_r = index.add %woc1, %state_stride : index + %woc2 = index.assume %woc2_r [range(%woc2_r, 0, 1073741820)] : index + vector.store %value2, %output[%woc2] : vector<4xf32>, view<1073741824xf32> + %woc3_r = index.add %woc2, %state_stride : index + %woc3 = index.assume %woc3_r [range(%woc3_r, 0, 1073741820)] : index + vector.store %value3, %output[%woc3] : vector<4xf32>, view<1073741824xf32> + %woc4_r = index.add %woc3, %state_stride : index + %woc4 = index.assume %woc4_r [range(%woc4_r, 0, 1073741820)] : index + vector.store %value4, %output[%woc4] : vector<4xf32>, view<1073741824xf32> + %woc5_r = index.add %woc4, %state_stride : index + %woc5 = index.assume %woc5_r [range(%woc5_r, 0, 1073741820)] : index + vector.store %value5, %output[%woc5] : vector<4xf32>, view<1073741824xf32> + %woc6_r = index.add %woc5, %state_stride : index + %woc6 = index.assume %woc6_r [range(%woc6_r, 0, 1073741820)] : index + vector.store %value6, %output[%woc6] : vector<4xf32>, view<1073741824xf32> + %woc7_r = index.add %woc6, %state_stride : index + %woc7 = index.assume %woc7_r [range(%woc7_r, 0, 1073741820)] : index + vector.store %value7, %output[%woc7] : vector<4xf32>, view<1073741824xf32> + %woc8_r = index.add %woc7, %state_stride : index + %woc8 = index.assume %woc8_r [range(%woc8_r, 0, 1073741820)] : index + vector.store %value8, %output[%woc8] : vector<4xf32>, view<1073741824xf32> + %woc9_r = index.add %woc8, %state_stride : index + %woc9 = index.assume %woc9_r [range(%woc9_r, 0, 1073741820)] : index + vector.store %value9, %output[%woc9] : vector<4xf32>, view<1073741824xf32> + %woc10_r = index.add %woc9, %state_stride : index + %woc10 = index.assume %woc10_r [range(%woc10_r, 0, 1073741820)] : index + vector.store %value10, %output[%woc10] : vector<4xf32>, view<1073741824xf32> + %woc11_r = index.add %woc10, %state_stride : index + %woc11 = index.assume %woc11_r [range(%woc11_r, 0, 1073741820)] : index + vector.store %value11, %output[%woc11] : vector<4xf32>, view<1073741824xf32> + %woc12_r = index.add %woc11, %state_stride : index + %woc12 = index.assume %woc12_r [range(%woc12_r, 0, 1073741820)] : index + vector.store %value12, %output[%woc12] : vector<4xf32>, view<1073741824xf32> + %woc13_r = index.add %woc12, %state_stride : index + %woc13 = index.assume %woc13_r [range(%woc13_r, 0, 1073741820)] : index + vector.store %value13, %output[%woc13] : vector<4xf32>, view<1073741824xf32> + %woc14_r = index.add %woc13, %state_stride : index + %woc14 = index.assume %woc14_r [range(%woc14_r, 0, 1073741820)] : index + vector.store %value14, %output[%woc14] : vector<4xf32>, view<1073741824xf32> + %woc15_r = index.add %woc14, %state_stride : index + %woc15 = index.assume %woc15_r [range(%woc15_r, 0, 1073741820)] : index + vector.store %value15, %output[%woc15] : vector<4xf32>, view<1073741824xf32> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/publish_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/publish_q8.loom new file mode 100644 index 000000000000..afa7376cbdf4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/publish_q8.loom @@ -0,0 +1,11 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @llm.gated_delta_net.publish_q8.vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.def<@llm.gated_delta_net.publish_q8.vector4_strict> device @llm_gated_delta_net_publish_q8_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + template.apply<@ggml.quantize_q8_1_x4.publish_vector4_strict>(%publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %output) : (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/gated_delta_net/rmsnorm_gate.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/rmsnorm_gate.loom new file mode 100644 index 000000000000..bc02b5fe0b72 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/rmsnorm_gate.loom @@ -0,0 +1,62 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.publish_q8.vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.gated_delta_net.rmsnorm_gate.publish_head128(%base: offset, %n_tokens: index, %n_heads: index, %n_seqs: index, %lane: index, %wave: index, %time_base: index, %seq: index, %zero: index, %head: index, %c4: index, %c8: index, %c16: index, %c128: index, %rms_epsilon: f32, %gate_op: index, %publish_q8: i1, %lds: buffer, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer) + +template.def<@llm.gated_delta_net.rmsnorm_gate.publish_head128> device @llm_gated_delta_net_rmsnorm_gate_publish_head128(%base: offset, %n_tokens: index, %n_heads: index, %n_seqs: index, %lane: index, %wave: index, %time_base: index, %seq: index, %zero: index, %head: index, %c4: index, %c8: index, %c16: index, %c128: index, %rms_epsilon: f32, %gate_op: index, %publish_q8: i1, %lds: buffer, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer) { + // One native subgroup handles each local attention row, using the + // standalone RMS reduction and unary implementation without a global read. + %post_rows0 = index.mul %n_tokens, %n_heads : index + %post_rows = index.mul %post_rows0, %n_seqs : index + %post_tokens = index.mul %n_tokens, %n_seqs : index + %post_groups = index.mul %n_heads, %c8 : index + %post_channel0 = index.mul %lane, %c4 : index + %post_channel = index.assume %post_channel0 [range(%post_channel0, 0, 124), mul(%post_channel0, 4)] : index + %post_attention = buffer.view %lds[%base] : buffer -> view<16x128xf32> + %post_weights = buffer.view %rms_weight[%base] : buffer -> view<128xf32> + %post_gates = buffer.view %raw_gate[%base] : buffer -> view<[%post_rows]x128xf32> + %post_f32 = buffer.view %norm_output[%base] : buffer -> view<[%post_rows]x128xf32> + %post_f16 = buffer.view %half_output[%base] : buffer -> view<[%post_groups]x[%post_tokens]x16xf16> + scf.for %post_local0 = [%wave to %c16 step %c8] { + %post_local = index.assume %post_local0 [range(%post_local0, 0, 15)] : index + %post_time0 = index.add %time_base, %post_local : index + %post_valid = index.cmp ult, %post_time0, %n_tokens : index + %post_time_safe = scf.select %post_valid, %post_time0, %zero : index + %post_time = index.assume %post_time_safe [lt(%post_time_safe, %n_tokens)] : index + %post_token0 = index.madd %seq, %n_tokens, %post_time : index + %post_token = index.assume %post_token0 [lt(%post_token0, %post_tokens)] : index + %post_row0 = index.madd %post_token, %n_heads, %head : index + %post_row = index.assume %post_row0 [lt(%post_row0, %post_rows)] : index + %post_scale = template.apply<@ggml.rmsnorm_f32.subgroup_row_scale>(%c16, %post_local, %c128, %rms_epsilon, %lds) : (index, index, index, f32, buffer) -> (f32) + scf.if %post_valid { + %post_x = vector.load %post_attention[%post_local, %post_channel] : view<16x128xf32> -> vector<4xf32> + %post_w = vector.load %post_weights[%post_channel] : view<128xf32> -> vector<4xf32> + %post_g = vector.load %post_gates[%post_row, %post_channel] : view<[%post_rows]x128xf32> -> vector<4xf32> + %post_scale_v = vector.splat %post_scale : vector<4xf32> + %post_normalized = vector.mulf %post_x, %post_scale_v : vector<4xf32> + %post_weighted = vector.mulf %post_normalized, %post_w : vector<4xf32> + %post_activated = template.apply<@ggml.unary_f32.apply_vector4>(%gate_op, %post_g) : (index, vector<4xf32>) -> (vector<4xf32>) + %post_result = vector.mulf %post_weighted, %post_activated : vector<4xf32> + vector.store %post_result, %post_f32[%post_row, %post_channel] : vector<4xf32>, view<[%post_rows]x128xf32> + scf.if %publish_q8 { + %post_flat_channel = index.madd %post_row, %c128, %post_channel : index + %post_scratch_values = buffer.view %lds[%base] : buffer -> view<256xf32> + %post_scratch_d = buffer.view %lds[%base] : buffer -> view<32xf32> + template.apply<@llm.gated_delta_net.publish_q8.vector4_strict>(%post_valid, %base, %post_flat_channel, %post_result, %post_scratch_values, %post_scratch_d, %half_output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } else { + %post_half = vector.fptrunc %post_result : vector<4xf32> to vector<4xf16> + %post_head_channel = index.madd %head, %c128, %post_channel : index + %post_group = index.div %post_head_channel, %c16 : index + %post_group_lane = index.rem %post_head_channel, %c16 : index + vector.store %post_half, %post_f16[%post_group, %post_token, %post_group_lane] : vector<4xf16>, view<[%post_groups]x[%post_tokens]x16xf16> + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom new file mode 100644 index 000000000000..93042dd9d061 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/scan_state_head128.loom @@ -0,0 +1,3330 @@ +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) +template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) +config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] + +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.snapshot_prefix.body(%state_factor: f32, %cache_base: index, %state_offset: index, %wcol: index, %row0: index, %rowh: index, %wave: index, %lane: index, %lds: buffer, %snapshot_cache: buffer, %sa0: vector<4xf32>, %sa1: vector<4xf32>, %sa2: vector<4xf32>, %sa3: vector<4xf32>, %sa4: vector<4xf32>, %sa5: vector<4xf32>, %sa6: vector<4xf32>, %sa7: vector<4xf32>, %sa8: vector<4xf32>, %sa9: vector<4xf32>, %sa10: vector<4xf32>, %sa11: vector<4xf32>, %sa12: vector<4xf32>, %sa13: vector<4xf32>, %sa14: vector<4xf32>, %sa15: vector<4xf32>, %fz8: vector<8xf32>) + +template.decl @llm.gated_delta_net.snapshot_fragment_prefix.body(%state_factor: f32, %cache_base: index, %state_offset: index, %wcol: index, %row0: index, %rowh: index, %wave: index, %lane: index, %lds: buffer, %snapshot_cache: buffer, %lo0_state: vector<8xf32>, %lo1_state: vector<8xf32>, %lo2_state: vector<8xf32>, %lo3_state: vector<8xf32>, %hi0_state: vector<8xf32>, %hi1_state: vector<8xf32>, %hi2_state: vector<8xf32>, %hi3_state: vector<8xf32>, %fz8: vector<8xf32>) + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +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.rmsnorm_gate.publish_head128(%base: offset, %n_tokens: index, %n_heads: index, %n_seqs: index, %lane: index, %wave: index, %time_base: index, %seq: index, %zero: index, %head: index, %c4: index, %c8: index, %c16: index, %c128: index, %rms_epsilon: f32, %gate_op: index, %publish_q8: i1, %lds: buffer, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer) + +template.decl @llm.gated_delta_net.inplace_state.publish_head128(%state: view<1073741824xf32>, %index0: index, %index1: index, %index2: index, %index3: index, %index4: index, %index5: index, %index6: index, %index7: index, %index8: index, %index9: index, %index10: index, %index11: index, %index12: index, %index13: index, %index14: index, %index15: index, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) + +template.decl @llm.gated_delta_net.snapshot_load.vector4(%state: view<1073741824xf32>, %index: index, %state_offset: index, %selected_state_offset: index, %selected_valid: i1, %zero: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.gated_delta_net.snapshot_publish.final_head128(%cache: view<1073741824xf32>, %index0: index, %index1: index, %index2: index, %index3: index, %index4: index, %index5: index, %index6: index, %index7: index, %index8: index, %index9: index, %index10: index, %index11: index, %index12: index, %index13: index, %index14: index, %index15: index, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) + +template.decl @llm.gated_delta_net.selected_snapshot_load.state_offset(%select_state: i1, %state_row_count: index, %state_ids: buffer, %base: offset, %sequence_count: index, %sequence: index, %head_count: index, %head: index, %state_size: index, %default_offset: index) -> (index, i1) + +template.decl @llm.gated_delta_net.publish_f32.scalar(%value: f32, %output: view<1073741824xf32>, %index: index) + +template.decl @llm.gated_delta_net.publish_f32.state_head128(%column: index, %head_width: index, %state_offset: index, %attention_elements: index, %row: index, %state_stride: index, %output: view<1073741824xf32>, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) + +// Chunked Gated DeltaNet matrix-core scan with one eight-wave workgroup per head and sequence. +// Each lane retains four rows of every state column in registers and stages the shard to F16 LDS once per 16-token chunk. +// The 63,296-byte LDS footprint permits one full-head workgroup per CU. +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.token_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.gated_delta_net.state_row_count : %value: index where [range(%value, 1, 262208)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.qk_stride1 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.qk_stride2 : %value: index where [range(%value, 0, 1073741823), mul(%value, 8)] + +config.decl @llm.gated_delta_net.qk_stride3 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.value_stride1 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.value_stride2 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.value_stride3 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.scalar_stride1 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.scalar_stride2 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.scalar_stride3 : %value: index where [range(%value, 0, 1073741823)] + +config.decl @llm.gated_delta_net.query_head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.query_sequence_ratio : %value: index where [range(%value, 1, 1024)] + +config.decl @llm.gated_delta_net.snapshot_stride : %value: index where [range(%value, 1, 1073741823)] + +config.decl @llm.gated_delta_net.l2_epsilon : f32 + +config.decl @llm.gated_delta_net.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] + +template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta_net_f32_wmma_head128_body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %k0 = index.constant 0 : index + %k1 = index.constant 1 : index + %k2 = index.constant 2 : index + %k3 = index.constant 3 : index + %k4 = index.constant 4 : index + %k5 = index.constant 5 : index + %k6 = index.constant 6 : index + %k7 = index.constant 7 : index + %k8 = index.constant 8 : index + %k9 = index.constant 9 : index + %k10 = index.constant 10 : index + %k11 = index.constant 11 : index + %k12 = index.constant 12 : index + %k13 = index.constant 13 : index + %k14 = index.constant 14 : index + %k15 = index.constant 15 : index + %k16 = index.constant 16 : index + %k17 = index.constant 17 : index + %k20 = index.constant 20 : index + %k24 = index.constant 24 : index + %k28 = index.constant 28 : index + %k32 = index.constant 32 : index + %k36 = index.constant 36 : index + %k40 = index.constant 40 : index + %k44 = index.constant 44 : index + %k48 = index.constant 48 : index + %k52 = index.constant 52 : index + %k56 = index.constant 56 : index + %k60 = index.constant 60 : index + %k61 = index.constant 61 : index + %k64 = index.constant 64 : index + %k65 = index.constant 65 : index + %k66 = index.constant 66 : index + %k67 = index.constant 67 : index + %k68 = index.constant 68 : index + %k69 = index.constant 69 : index + %k70 = index.constant 70 : index + %k71 = index.constant 71 : index + %k72 = index.constant 72 : index + %k73 = index.constant 73 : index + %k74 = index.constant 74 : index + %k75 = index.constant 75 : index + %k76 = index.constant 76 : index + %k77 = index.constant 77 : index + %k78 = index.constant 78 : index + %k79 = index.constant 79 : index + %k80 = index.constant 80 : index + %k88 = index.constant 88 : index + %k96 = index.constant 96 : index + %k104 = index.constant 104 : index + %k112 = index.constant 112 : index + %k120 = index.constant 120 : index + %k128 = index.constant 128 : index + %k136 = index.constant 136 : index + %k144 = index.constant 144 : index + %k152 = index.constant 152 : index + %k160 = index.constant 160 : index + %k168 = index.constant 168 : index + %k176 = index.constant 176 : index + %k184 = index.constant 184 : index + %k192 = index.constant 192 : index + %k200 = index.constant 200 : index + %k208 = index.constant 208 : index + %k216 = index.constant 216 : index + %k224 = index.constant 224 : index + %k232 = index.constant 232 : index + %k240 = index.constant 240 : index + %k248 = index.constant 248 : index + %k256 = index.constant 256 : index + %k272 = index.constant 272 : index + %k384 = index.constant 384 : index + %k408 = index.constant 408 : index + %k512 = index.constant 512 : index + %k544 = index.constant 544 : index + %k640 = index.constant 640 : index + %k680 = index.constant 680 : index + %k768 = index.constant 768 : index + %k816 = index.constant 816 : index + %k896 = index.constant 896 : index + %k952 = index.constant 952 : index + %k1024 = index.constant 1024 : index + %k1088 = index.constant 1088 : index + %k1152 = index.constant 1152 : index + %k1224 = index.constant 1224 : index + %k1280 = index.constant 1280 : index + %k1360 = index.constant 1360 : index + %k1408 = index.constant 1408 : index + %k1496 = index.constant 1496 : index + %k1536 = index.constant 1536 : index + %k1632 = index.constant 1632 : index + %k1664 = index.constant 1664 : index + %k1768 = index.constant 1768 : index + %k1792 = index.constant 1792 : index + %k1904 = index.constant 1904 : index + %k1920 = index.constant 1920 : index + %k2040 = index.constant 2040 : index + %k4352 = index.constant 4352 : index + %k5440 = index.constant 5440 : index + %k6528 = index.constant 6528 : index + %k11072 = index.constant 11072 : index + + %s_v0 = config.get @llm.gated_delta_net.head_width : index + %n_heads0 = config.get @llm.gated_delta_net.head_count : index + %n_tokens0 = config.get @llm.gated_delta_net.token_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %sq1 = config.get @llm.gated_delta_net.qk_stride1 : index + %sq2 = config.get @llm.gated_delta_net.qk_stride2 : index + %sq3 = config.get @llm.gated_delta_net.qk_stride3 : index + %sv1 = config.get @llm.gated_delta_net.value_stride1 : index + %sv2 = config.get @llm.gated_delta_net.value_stride2 : index + %sv3 = config.get @llm.gated_delta_net.value_stride3 : index + %sb1 = config.get @llm.gated_delta_net.scalar_stride1 : index + %sb2 = config.get @llm.gated_delta_net.scalar_stride2 : index + %sb3 = config.get @llm.gated_delta_net.scalar_stride3 : index + %neqk1 = config.get @llm.gated_delta_net.query_head_count : index + %rq3 = config.get @llm.gated_delta_net.query_sequence_ratio : index + %s_v = index.assume %s_v0 [range(%s_v0, 128, 128)] : index + %n_heads = index.assume %n_heads0 [range(%n_heads0, 1, 4096)] : index + %n_tokens = index.assume %n_tokens0 [range(%n_tokens0, 1, 1048576)] : index + + %recurrence_true = scalar.constant true : i1 + %recurrence_wide = index.cmp ugt, %n_tokens, %k5 : index + %recurrence_other = scalar.xori %publish_q8, %recurrence_true : i1 + %all_recurrence_rows = scalar.ori %recurrence_wide, %recurrence_other : i1 + + %h_idx0 = kernel.workgroup.id : index + %h_idx = index.assume %h_idx0 [range(%h_idx0, 0, 4095)] : index + %seq0 = kernel.workgroup.id : index + %seq = index.assume %seq0 [range(%seq0, 0, 4095)] : index + %colblk0 = kernel.workgroup.id : index + %colblk = index.assume %colblk0 [range(%colblk0, 0, 1023)] : index + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index + %wave0 = kernel.subgroup.id : index + %wave = index.assume %wave0 [range(%wave0, 0, 7)] : index + %is_w0 = index.cmp eq, %wave, %c0 : index + %is_w1 = index.cmp eq, %wave, %c1 : index + %lane_live = index.cmp ult, %lane, %k16 : index + %lane_lo = index.cmp ult, %lane, %k16 : index + + %bb = index.mul %colblk, %k128 : index + %wo = index.mul %wave, %k16 : index + %waveoff0 = index.assume %wo [range(%wo, 0, 112)] : index + %wcol0 = index.add %bb, %wo : index + %wcol = index.assume %wcol0 [range(%wcol0, 0, 127)] : index + + %iq1 = index.rem %h_idx, %neqk1 : index + %iq3 = index.div %seq, %rq3 : index + + %q_g = buffer.assume.memory_space %q : buffer + %k_g = buffer.assume.memory_space %k : buffer + %v_g = buffer.assume.memory_space %v : buffer + %g_g = buffer.assume.memory_space %g : buffer + %b_g = buffer.assume.memory_space %beta : buffer + %bias_g = buffer.assume.memory_space %bias : buffer + %a_scale_g = buffer.assume.memory_space %a_scale : buffer + %si_g = buffer.assume.memory_space %state_in : buffer + %d_g = buffer.assume.memory_space %dst : buffer + %snapshot_cache_g = buffer.assume.memory_space %snapshot_cache : buffer + %q_view = buffer.view %q_g[%base] : buffer -> view<1073741824xf32> + %k_view = buffer.view %k_g[%base] : buffer -> view<1073741824xf32> + %v_view = buffer.view %v_g[%base] : buffer -> view<1073741824xf32> + %g_view = buffer.view %g_g[%base] : buffer -> view<1073741824xf32> + %b_view = buffer.view %b_g[%base] : buffer -> view<1073741824xf32> + %bias_view = buffer.view %bias_g[%base] : buffer -> view<4096xf32> + %a_scale_view = buffer.view %a_scale_g[%base] : buffer -> view<4096xf32> + %si_view = buffer.view %si_g[%base] : buffer -> view<1073741824xf32> + %dst_view = buffer.view %d_g[%base] : buffer -> view<1073741824xf32> + %snapshot_cache_view = buffer.view %snapshot_cache_g[%base] : buffer -> view<1073741824xf32> + + %ae0 = index.mul %s_v, %n_heads : index + %ae1 = index.mul %ae0, %n_tokens : index + %attn_elems = index.mul %ae1, %n_seqs : index + %sv_sq = index.mul %s_v, %s_v : index + %so0 = index.mul %seq, %n_heads : index + %so1 = index.add %so0, %h_idx : index + %state_offset = index.mul %so1, %sv_sq : index + %selected_state_offset, %selected_valid = template.apply<@llm.gated_delta_net.selected_snapshot_load.state_offset>(%select_state, %state_row_count, %state_ids, %base, %n_seqs, %seq, %n_heads, %h_idx, %sv_sq, %state_offset) : (i1, index, buffer, offset, index, index, index, index, index, index) -> (index, i1) + %selected_zero_state = vector.constant 0.0 : vector<4xf32> + %ab0 = index.mul %seq, %n_tokens : index + %ab1 = index.mul %ab0, %n_heads : index + %ab2 = index.add %ab1, %h_idx : index + %attn_base = index.mul %ab2, %s_v : index + %attn_stride = index.mul %s_v, %n_heads : index + %qk_s = index.mul %iq3, %sq3 : index + %qk_h = index.mul %iq1, %sq1 : index + %qk_base = index.add %qk_s, %qk_h : index + %v_s = index.mul %seq, %sv3 : index + %v_h = index.mul %h_idx, %sv1 : index + %v_head_base = index.add %v_s, %v_h : index + %gb_s = index.mul %seq, %sb3 : index + %gb_h = index.mul %h_idx, %sb1 : index + %gb_r = index.add %gb_s, %gb_h : index + %gb_base = index.assume %gb_r [range(%gb_r, 0, 1073741823)] : index + + %lds_bytes = index.constant 63296 : offset + %lds = buffer.alloca align(16) %lds_bytes : buffer + %km_o = index.constant 0 : offset + %qm_o = index.constant 4352 : offset + %ka_o = index.constant 4352 : offset + %gr_o = index.constant 8704 : offset + %qkt_o = index.constant 9728 : offset + %lay_tokrow = encoding.layout.strided [136, 1] : encoding + %lay_rowtok = encoding.layout.strided [1, 136] : encoding + %lay_tile = encoding.layout.strided [17, 1] : encoding + %lay_gram = encoding.layout.strided [16, 1] : encoding + %lay_colrow = encoding.layout.strided [136, 1] : encoding + %lay_dcol = encoding.layout.strided [24, 1] : encoding + %lay_katok = encoding.layout.strided [1, 24] : encoding + %lay_upd = encoding.layout.strided [68, 1] : encoding + %km_flat = buffer.view %lds[%km_o] : buffer -> view<2176xf16> + %qm_flat = buffer.view %lds[%qm_o] : buffer -> view<2176xf16> + %ka_flat = buffer.view %lds[%ka_o] : buffer -> view<3072xf16> + %gr_flat = buffer.view %lds[%gr_o] : buffer -> view<256xf32> + %qkt_flat = buffer.view %lds[%qkt_o] : buffer -> view<256xf32> + %km_lhs = buffer.view %lds[%km_o] : buffer -> view<16x128xf16, %lay_tokrow> + %km_rhs = buffer.view %lds[%km_o] : buffer -> view<128x16xf16, %lay_rowtok> + %qm_rhs = buffer.view %lds[%qm_o] : buffer -> view<128x16xf16, %lay_rowtok> + %ka_rhs = buffer.view %lds[%ka_o] : buffer -> view<16x128xf16, %lay_katok> + %gr_res = buffer.view %lds[%gr_o] : buffer -> view<16x16xf32, %lay_gram> + %qkt_res = buffer.view %lds[%qkt_o] : buffer -> view<16x16xf32, %lay_gram> + + %wb0 = index.mul %wave, %k6528 : index + %wb1 = index.add %wb0, %k11072 : index + %wb = index.cast %wb1 : index to offset + %ab1o = index.add %wb1, %k4352 : index + %abo = index.cast %ab1o : index to offset + %bb1o = index.add %wb1, %k5440 : index + %bbo = index.cast %bb1o : index to offset + %dd1o = index.add %wb1, %k4352 : index + %ddo = index.cast %dd1o : index to offset + %sh_flat = buffer.view %lds[%wb] : buffer -> view<2176xf16> + %sh_lhs = buffer.view %lds[%wb] : buffer -> view<16x128xf16, %lay_colrow> + %upd_flat = buffer.view %lds[%wb] : buffer -> view<1088xf32> + %upd_tile = buffer.view %lds[%wb] : buffer -> view<16x68xf32, %lay_upd> + %a_flat = buffer.view %lds[%abo] : buffer -> view<272xf32> + %a_res = buffer.view %lds[%abo] : buffer -> view<16x16xf32, %lay_tile> + %b_flat = buffer.view %lds[%bbo] : buffer -> view<272xf32> + %b_res = buffer.view %lds[%bbo] : buffer -> view<16x16xf32, %lay_tile> + %d_flat = buffer.view %lds[%ddo] : buffer -> view<384xf16> + %d_lhs = buffer.view %lds[%ddo] : buffer -> view<16x16xf16, %lay_dcol> + // V and attention output share LDS slots; each recurrence reads V before overwriting its dead slot. + // The full-head plane overlays dead K/Q storage after the matrix products. + %vst_flat = buffer.view %lds[%km_o] : buffer -> view<2048xf32> + %aost_flat = buffer.view %lds[%km_o] : buffer -> view<2048xf32> + %vs_tok0 = index.div %tid, %k128 : index + %vs_tok = index.assume %vs_tok0 [range(%vs_tok0, 0, 1)] : index + %vs_col0 = index.rem %tid, %k128 : index + %vs_col = index.assume %vs_col0 [range(%vs_col0, 0, 127)] : index + %vs_gcol0 = index.add %bb, %vs_col : index + %vs_gcol = index.assume %vs_gcol0 [range(%vs_gcol0, 0, 127)] : index + %gs_o = index.constant 10752 : offset + %gs_flat = buffer.view %lds[%gs_o] : buffer -> view<80xf32> + + // Low-row projection epilogue. Distribute the valid token chains across + // waves and issue them before the state-plane loads below so the scheduler + // can cover their transcendental latency with independent global reads. + scf.if %projection_epilogue { + %projection_zero = scalar.constant 0.0 : f32 + %projection_one = scalar.constant 1.0 : f32 + %projection_worker = index.cmp ult, %lane, %k2 : index + %projection_lane_offset = index.mul %lane, %k8 : index + %projection_slot_r = index.add %projection_lane_offset, %wave : index + %projection_slot = index.assume %projection_slot_r [range(%projection_slot_r, 0, 15)] : index + %projection_token_ok = index.cmp ult, %projection_slot, %n_tokens : index + scf.if %projection_worker { + %gate_value, %beta_value = scf.if %projection_token_ok -> (f32, f32) { + %projection_token = index.assume %projection_slot [range(%projection_slot, 0, 1048575)] : index + %gbi_m = index.mul %projection_token, %sb2 : index + %gbi_a = index.add %gb_base, %gbi_m : index + %gbi = index.assume %gbi_a [range(%gbi_a, 0, 1073741823)] : index + %graw = view.load %g_view[%gbi] : view<1073741824xf32> -> f32 + %braw = view.load %b_view[%gbi] : view<1073741824xf32> -> f32 + %bias_value = view.load %bias_view[%h_idx] : view<4096xf32> -> f32 + %a_value = view.load %a_scale_view[%h_idx] : view<4096xf32> -> f32 + %gate_result, %beta_result = template.apply<@llm.gated_delta_net.projection_epilogue.apply>(%graw, %braw, %bias_value, %a_value, %projection_zero, %projection_one) : (f32, f32, f32, f32, f32, f32) -> (f32, f32) + scf.yield %gate_result, %beta_result : f32, f32 + } else { + scf.yield %projection_zero, %projection_zero : f32, f32 + } + %gsi = index.add %projection_slot, %k64 : index + %gsib = index.assume %gsi [range(%gsi, 64, 79)] : index + view.store %gate_value, %gs_flat[%gsib] : f32, view<80xf32> + %bsi_m = index.mul %projection_slot, %k4 : index + %bsib = index.assume %bsi_m [range(%bsi_m, 0, 60)] : index + view.store %beta_value, %gs_flat[%bsib] : f32, view<80xf32> + } + } + %projection_true = scalar.constant true : i1 + %load_processed_scalars = scalar.xori %projection_epilogue, %projection_true : i1 + + // Staging map: 16 consecutive lanes cover one token's 128 rows, 8 rows each, + // so a pass reads whole 512 B token rows and the wave's accesses coalesce. + %tid_lt_c = index.cmp ult, %tid, %k16 : index + %stokl_r = index.div %tid, %k16 : index + %stokl = index.assume %stokl_r [range(%stokl_r, 0, 15)] : index + %srlane = index.rem %tid, %k16 : index + %srow0_m = index.mul %srlane, %k8 : index + %srow0 = index.assume %srow0_m [range(%srow0_m, 0, 120)] : index + + // Lane L owns rows 4L..4L+3 of every column: one f32x4 per column, and the + // update's two 64-row halves split exactly on the lane-16 boundary. + %row0_m = index.mul %lane, %k4 : index + %row0 = index.assume %row0_m [range(%row0_m, 0, 124)] : index + %lc_s = scf.select %lane_live, %lane, %c0 : index + %lc = index.assume %lc_s [range(%lc_s, 0, 15)] : index + %rec_col0 = index.add %waveoff0, %lc : index + %rec_col = index.assume %rec_col0 [range(%rec_col0, 0, 127)] : index + %gcol_r = index.add %wcol, %lc : index + %gcol = index.assume %gcol_r [range(%gcol_r, 0, 127)] : index + %tile_row_m = index.mul %lc, %k17 : index + %tile_row = index.assume %tile_row_m [range(%tile_row_m, 0, 255)] : index + %d_row_m = index.mul %lc, %k24 : index + %d_row = index.assume %d_row_m [range(%d_row_m, 0, 360)] : index + %ic0_r = index.add %wcol, %k0 : index + %ic0 = index.assume %ic0_r [range(%ic0_r, 0, 127)] : index + %icb0_m = index.mul %ic0, %s_v : index + %icb0_a = index.add %state_offset, %icb0_m : index + %icb0_r = index.add %icb0_a, %row0 : index + %icb0 = index.assume %icb0_r [range(%icb0_r, 0, 1073741820)] : index + %s0_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb0, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb1_r = index.add %icb0, %k128 : index + %icb1 = index.assume %icb1_r [range(%icb1_r, 0, 1073741820)] : index + %s1_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb1, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb2_r = index.add %icb1, %k128 : index + %icb2 = index.assume %icb2_r [range(%icb2_r, 0, 1073741820)] : index + %s2_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb2, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb3_r = index.add %icb2, %k128 : index + %icb3 = index.assume %icb3_r [range(%icb3_r, 0, 1073741820)] : index + %s3_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb3, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb4_r = index.add %icb3, %k128 : index + %icb4 = index.assume %icb4_r [range(%icb4_r, 0, 1073741820)] : index + %s4_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb4, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb5_r = index.add %icb4, %k128 : index + %icb5 = index.assume %icb5_r [range(%icb5_r, 0, 1073741820)] : index + %s5_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb5, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb6_r = index.add %icb5, %k128 : index + %icb6 = index.assume %icb6_r [range(%icb6_r, 0, 1073741820)] : index + %s6_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb6, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb7_r = index.add %icb6, %k128 : index + %icb7 = index.assume %icb7_r [range(%icb7_r, 0, 1073741820)] : index + %s7_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb7, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb8_r = index.add %icb7, %k128 : index + %icb8 = index.assume %icb8_r [range(%icb8_r, 0, 1073741820)] : index + %s8_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb8, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb9_r = index.add %icb8, %k128 : index + %icb9 = index.assume %icb9_r [range(%icb9_r, 0, 1073741820)] : index + %s9_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb9, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb10_r = index.add %icb9, %k128 : index + %icb10 = index.assume %icb10_r [range(%icb10_r, 0, 1073741820)] : index + %s10_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb10, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb11_r = index.add %icb10, %k128 : index + %icb11 = index.assume %icb11_r [range(%icb11_r, 0, 1073741820)] : index + %s11_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb11, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb12_r = index.add %icb11, %k128 : index + %icb12 = index.assume %icb12_r [range(%icb12_r, 0, 1073741820)] : index + %s12_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb12, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb13_r = index.add %icb12, %k128 : index + %icb13 = index.assume %icb13_r [range(%icb13_r, 0, 1073741820)] : index + %s13_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb13, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb14_r = index.add %icb13, %k128 : index + %icb14 = index.assume %icb14_r [range(%icb14_r, 0, 1073741820)] : index + %s14_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb14, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %icb15_r = index.add %icb14, %k128 : index + %icb15 = index.assume %icb15_r [range(%icb15_r, 0, 1073741820)] : index + %s15_init = template.apply<@llm.gated_delta_net.snapshot_load.vector4>(%si_view, %icb15, %state_offset, %selected_state_offset, %selected_valid, %selected_zero_state) : (view<1073741824xf32>, index, index, index, i1, vector<4xf32>) -> (vector<4xf32>) + %ch_r = index.add %n_tokens, %k15 : index + %n_chunks = index.div %ch_r, %k16 : index + %f1 = scalar.constant 1.0 : f32 + %f0 = scalar.constant 0.0 : f32 + %norm_eps = config.get @llm.gated_delta_net.l2_epsilon : f32 + %norm_eps_sq = scalar.mulf %norm_eps, %norm_eps : f32 + %fz8 = vector.constant 0.0 : vector<8xf32> + %zv8 = vector.constant 0.0 : vector<8xf32> + %fzh8 = vector.constant 0.0 : vector<8xf16> + %five_snapshot_tokens = index.cmp eq, %n_tokens, %k5 : index + %reuse_snapshot_fragments = scalar.andi %publish_snapshots, %five_snapshot_tokens : i1 + %sf0, %sf1, %sf2, %sf3, %sf4, %sf5, %sf6, %sf7, %sf8, %sf9, %sf10, %sf11, %sf12, %sf13, %sf14, %sf15 = scf.for %ch = [%c0 to %n_chunks step %c1](%sa0 = %s0_init : vector<4xf32>, %sa1 = %s1_init : vector<4xf32>, %sa2 = %s2_init : vector<4xf32>, %sa3 = %s3_init : vector<4xf32>, %sa4 = %s4_init : vector<4xf32>, %sa5 = %s5_init : vector<4xf32>, %sa6 = %s6_init : vector<4xf32>, %sa7 = %s7_init : vector<4xf32>, %sa8 = %s8_init : vector<4xf32>, %sa9 = %s9_init : vector<4xf32>, %sa10 = %s10_init : vector<4xf32>, %sa11 = %s11_init : vector<4xf32>, %sa12 = %s12_init : vector<4xf32>, %sa13 = %s13_init : vector<4xf32>, %sa14 = %s14_init : vector<4xf32>, %sa15 = %s15_init : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %t0 = index.mul %ch, %k16 : index + // Compute gate prefix and suffix decays once in log space to avoid product underflow. + scf.if %load_processed_scalars { + // Masked tokens use g=1 and beta=0; threads 0..15 each own one token. + %gtok_r = index.add %t0, %tid : index + %gtok_ok = index.cmp ult, %gtok_r, %n_tokens : index + %gtok_s = scf.select %gtok_ok, %gtok_r, %c0 : index + %gtok = index.assume %gtok_s [range(%gtok_s, 0, 1048575)] : index + scf.if %tid_lt_c { + %gbi_m = index.mul %gtok, %sb2 : index + %gbi_a = index.add %gb_base, %gbi_m : index + %gbi = index.assume %gbi_a [range(%gbi_a, 0, 1073741823)] : index + %graw = view.load %g_view[%gbi] : view<1073741824xf32> -> f32 + %braw = view.load %b_view[%gbi] : view<1073741824xf32> -> f32 + %gvv = scf.select %gtok_ok, %graw, %f0 : f32 + %bvv = scf.select %gtok_ok, %braw, %f0 : f32 + %gsi = index.add %tid, %k64 : index + %gsib = index.assume %gsi [range(%gsi, 0, 79)] : index + view.store %gvv, %gs_flat[%gsib] : f32, view<80xf32> + %bsi_m = index.mul %tid, %k4 : index + %bsib = index.assume %bsi_m [range(%bsi_m, 0, 79)] : index + view.store %bvv, %gs_flat[%bsib] : f32, view<80xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // c_t, the suffix decay, and log(c_t). Every lane reads all 16 + // log-gates and forms prefix and suffix sums without a cross-lane scan. + scf.if %tid_lt_c { + %pg0 = view.load %gs_flat[%k64] : view<80xf32> -> f32 + %pg1 = view.load %gs_flat[%k65] : view<80xf32> -> f32 + %pg2 = view.load %gs_flat[%k66] : view<80xf32> -> f32 + %pg3 = view.load %gs_flat[%k67] : view<80xf32> -> f32 + %pg4 = view.load %gs_flat[%k68] : view<80xf32> -> f32 + %pg5 = view.load %gs_flat[%k69] : view<80xf32> -> f32 + %pg6 = view.load %gs_flat[%k70] : view<80xf32> -> f32 + %pg7 = view.load %gs_flat[%k71] : view<80xf32> -> f32 + %pg8 = view.load %gs_flat[%k72] : view<80xf32> -> f32 + %pg9 = view.load %gs_flat[%k73] : view<80xf32> -> f32 + %pg10 = view.load %gs_flat[%k74] : view<80xf32> -> f32 + %pg11 = view.load %gs_flat[%k75] : view<80xf32> -> f32 + %pg12 = view.load %gs_flat[%k76] : view<80xf32> -> f32 + %pg13 = view.load %gs_flat[%k77] : view<80xf32> -> f32 + %pg14 = view.load %gs_flat[%k78] : view<80xf32> -> f32 + %pg15 = view.load %gs_flat[%k79] : view<80xf32> -> f32 + %ple0 = index.cmp ule, %k0, %tid : index + %pm0 = scf.select %ple0, %pg0, %f0 : f32 + %sm0 = scf.select %ple0, %f0, %pg0 : f32 + %ple1 = index.cmp ule, %k1, %tid : index + %pm1 = scf.select %ple1, %pg1, %f0 : f32 + %sm1 = scf.select %ple1, %f0, %pg1 : f32 + %pp1 = scalar.addf %pm0, %pm1 : f32 + %ss1 = scalar.addf %sm0, %sm1 : f32 + %ple2 = index.cmp ule, %k2, %tid : index + %pm2 = scf.select %ple2, %pg2, %f0 : f32 + %sm2 = scf.select %ple2, %f0, %pg2 : f32 + %pp2 = scalar.addf %pp1, %pm2 : f32 + %ss2 = scalar.addf %ss1, %sm2 : f32 + %ple3 = index.cmp ule, %k3, %tid : index + %pm3 = scf.select %ple3, %pg3, %f0 : f32 + %sm3 = scf.select %ple3, %f0, %pg3 : f32 + %pp3 = scalar.addf %pp2, %pm3 : f32 + %ss3 = scalar.addf %ss2, %sm3 : f32 + %ple4 = index.cmp ule, %k4, %tid : index + %pm4 = scf.select %ple4, %pg4, %f0 : f32 + %sm4 = scf.select %ple4, %f0, %pg4 : f32 + %pp4 = scalar.addf %pp3, %pm4 : f32 + %ss4 = scalar.addf %ss3, %sm4 : f32 + %ple5 = index.cmp ule, %k5, %tid : index + %pm5 = scf.select %ple5, %pg5, %f0 : f32 + %sm5 = scf.select %ple5, %f0, %pg5 : f32 + %pp5 = scalar.addf %pp4, %pm5 : f32 + %ss5 = scalar.addf %ss4, %sm5 : f32 + %ple6 = index.cmp ule, %k6, %tid : index + %pm6 = scf.select %ple6, %pg6, %f0 : f32 + %sm6 = scf.select %ple6, %f0, %pg6 : f32 + %pp6 = scalar.addf %pp5, %pm6 : f32 + %ss6 = scalar.addf %ss5, %sm6 : f32 + %ple7 = index.cmp ule, %k7, %tid : index + %pm7 = scf.select %ple7, %pg7, %f0 : f32 + %sm7 = scf.select %ple7, %f0, %pg7 : f32 + %pp7 = scalar.addf %pp6, %pm7 : f32 + %ss7 = scalar.addf %ss6, %sm7 : f32 + %ple8 = index.cmp ule, %k8, %tid : index + %pm8 = scf.select %ple8, %pg8, %f0 : f32 + %sm8 = scf.select %ple8, %f0, %pg8 : f32 + %pp8 = scalar.addf %pp7, %pm8 : f32 + %ss8 = scalar.addf %ss7, %sm8 : f32 + %ple9 = index.cmp ule, %k9, %tid : index + %pm9 = scf.select %ple9, %pg9, %f0 : f32 + %sm9 = scf.select %ple9, %f0, %pg9 : f32 + %pp9 = scalar.addf %pp8, %pm9 : f32 + %ss9 = scalar.addf %ss8, %sm9 : f32 + %ple10 = index.cmp ule, %k10, %tid : index + %pm10 = scf.select %ple10, %pg10, %f0 : f32 + %sm10 = scf.select %ple10, %f0, %pg10 : f32 + %pp10 = scalar.addf %pp9, %pm10 : f32 + %ss10 = scalar.addf %ss9, %sm10 : f32 + %ple11 = index.cmp ule, %k11, %tid : index + %pm11 = scf.select %ple11, %pg11, %f0 : f32 + %sm11 = scf.select %ple11, %f0, %pg11 : f32 + %pp11 = scalar.addf %pp10, %pm11 : f32 + %ss11 = scalar.addf %ss10, %sm11 : f32 + %ple12 = index.cmp ule, %k12, %tid : index + %pm12 = scf.select %ple12, %pg12, %f0 : f32 + %sm12 = scf.select %ple12, %f0, %pg12 : f32 + %pp12 = scalar.addf %pp11, %pm12 : f32 + %ss12 = scalar.addf %ss11, %sm12 : f32 + %ple13 = index.cmp ule, %k13, %tid : index + %pm13 = scf.select %ple13, %pg13, %f0 : f32 + %sm13 = scf.select %ple13, %f0, %pg13 : f32 + %pp13 = scalar.addf %pp12, %pm13 : f32 + %ss13 = scalar.addf %ss12, %sm13 : f32 + %ple14 = index.cmp ule, %k14, %tid : index + %pm14 = scf.select %ple14, %pg14, %f0 : f32 + %sm14 = scf.select %ple14, %f0, %pg14 : f32 + %pp14 = scalar.addf %pp13, %pm14 : f32 + %ss14 = scalar.addf %ss13, %sm14 : f32 + %ple15 = index.cmp ule, %k15, %tid : index + %pm15 = scf.select %ple15, %pg15, %f0 : f32 + %sm15 = scf.select %ple15, %f0, %pg15 : f32 + %pp15 = scalar.addf %pp14, %pm15 : f32 + %ss15 = scalar.addf %ss14, %sm15 : f32 + %cpv = scalar.expf %pp15 : f32 + %pev = scalar.expf %ss15 : f32 + %quad = index.mul %tid, %k4 : index + %cpi = index.add %quad, %k1 : index + %cpib = index.assume %cpi [range(%cpi, 0, 79)] : index + view.store %cpv, %gs_flat[%cpib] : f32, view<80xf32> + %evi = index.add %quad, %k2 : index + %evib = index.assume %evi [range(%evi, 0, 79)] : index + view.store %pev, %gs_flat[%evib] : f32, view<80xf32> + %ivi = index.add %quad, %k3 : index + %ivib = index.assume %ivi [range(%ivi, 0, 79)] : index + view.store %pp15, %gs_flat[%ivib] : f32, view<80xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + // Normalize K/Q into F16 LDS while four-lane groups reproduce the exact wave32 reduction DAG. + %stk0_r = index.add %t0, %k0 : index + %stk0_s = index.add %stk0_r, %stokl : index + %stk0_ok = index.cmp ult, %stk0_s, %n_tokens : index + %stk0_c = scf.select %stk0_ok, %stk0_s, %c0 : index + %stk0 = index.assume %stk0_c [range(%stk0_c, 0, 1048575)] : index + %sgi0_m = index.mul %stk0, %sq2 : index + %sgi0_a = index.add %qk_base, %sgi0_m : index + %sgi0_b = index.add %sgi0_a, %srow0 : index + %sgi0 = index.assume %sgi0_b [range(%sgi0_b, 0, 1073741816)] : index + %kvr0 = vector.load %k_view[%sgi0] : view<1073741824xf32> -> vector<8xf32> + %kvs0 = scf.select %stk0_ok, %kvr0, %zv8 : vector<8xf32> + %n0_k_v0 = vector.extract %kvs0[0] : vector<8xf32> -> f32 + %n0_k_sq0 = scalar.mulf %n0_k_v0, %n0_k_v0 : f32 + %n0_k_v1 = vector.extract %kvs0[1] : vector<8xf32> -> f32 + %n0_k_sq1 = scalar.mulf %n0_k_v1, %n0_k_v1 : f32 + %n0_k_v2 = vector.extract %kvs0[2] : vector<8xf32> -> f32 + %n0_k_sq2 = scalar.mulf %n0_k_v2, %n0_k_v2 : f32 + %n0_k_v3 = vector.extract %kvs0[3] : vector<8xf32> -> f32 + %n0_k_sq3 = scalar.mulf %n0_k_v3, %n0_k_v3 : f32 + %n0_k_v4 = vector.extract %kvs0[4] : vector<8xf32> -> f32 + %n0_k_sq4 = scalar.mulf %n0_k_v4, %n0_k_v4 : f32 + %n0_k_v5 = vector.extract %kvs0[5] : vector<8xf32> -> f32 + %n0_k_sq5 = scalar.mulf %n0_k_v5, %n0_k_v5 : f32 + %n0_k_v6 = vector.extract %kvs0[6] : vector<8xf32> -> f32 + %n0_k_sq6 = scalar.mulf %n0_k_v6, %n0_k_v6 : f32 + %n0_k_v7 = vector.extract %kvs0[7] : vector<8xf32> -> f32 + %n0_k_sq7 = scalar.mulf %n0_k_v7, %n0_k_v7 : f32 + %n0_k_p01 = scalar.addf %n0_k_sq0, %n0_k_sq1 : f32 + %n0_k_p23 = scalar.addf %n0_k_sq2, %n0_k_sq3 : f32 + %n0_k_p45 = scalar.addf %n0_k_sq4, %n0_k_sq5 : f32 + %n0_k_p67 = scalar.addf %n0_k_sq6, %n0_k_sq7 : f32 + %n0_k_q03 = scalar.addf %n0_k_p01, %n0_k_p23 : f32 + %n0_k_q47 = scalar.addf %n0_k_p45, %n0_k_p67 : f32 + %n0_k_local8 = scalar.addf %n0_k_q03, %n0_k_q47 : f32 + %n0_k_c0_ge = index.cmp uge, %lane, %k0 : index + %n0_k_c0_lt = index.cmp ult, %lane, %k4 : index + %n0_k_c0_ge_v = scf.select %n0_k_c0_ge, %n0_k_local8, %f0 : f32 + %n0_k_c0_in = scf.select %n0_k_c0_lt, %n0_k_c0_ge_v, %f0 : f32 + %n0_k_w0 = kernel.subgroup.reduce %n0_k_c0_in : f32 + %n0_k_c2_ge = index.cmp uge, %lane, %k8 : index + %n0_k_c2_lt = index.cmp ult, %lane, %k12 : index + %n0_k_c2_ge_v = scf.select %n0_k_c2_ge, %n0_k_local8, %f0 : f32 + %n0_k_c2_in = scf.select %n0_k_c2_lt, %n0_k_c2_ge_v, %f0 : f32 + %n0_k_w2 = kernel.subgroup.reduce %n0_k_c2_in : f32 + %n0_k_c1_ge = index.cmp uge, %lane, %k4 : index + %n0_k_c1_lt = index.cmp ult, %lane, %k8 : index + %n0_k_c1_ge_v = scf.select %n0_k_c1_ge, %n0_k_local8, %f0 : f32 + %n0_k_c1_in = scf.select %n0_k_c1_lt, %n0_k_c1_ge_v, %f0 : f32 + %n0_k_w1 = kernel.subgroup.reduce %n0_k_c1_in : f32 + %n0_k_c3_ge = index.cmp uge, %lane, %k12 : index + %n0_k_c3_lt = index.cmp ult, %lane, %k16 : index + %n0_k_c3_ge_v = scf.select %n0_k_c3_ge, %n0_k_local8, %f0 : f32 + %n0_k_c3_in = scf.select %n0_k_c3_lt, %n0_k_c3_ge_v, %f0 : f32 + %n0_k_w3 = kernel.subgroup.reduce %n0_k_c3_in : f32 + %n0_k_c4_ge = index.cmp uge, %lane, %k16 : index + %n0_k_c4_lt = index.cmp ult, %lane, %k20 : index + %n0_k_c4_ge_v = scf.select %n0_k_c4_ge, %n0_k_local8, %f0 : f32 + %n0_k_c4_in = scf.select %n0_k_c4_lt, %n0_k_c4_ge_v, %f0 : f32 + %n0_k_w4 = kernel.subgroup.reduce %n0_k_c4_in : f32 + %n0_k_c6_ge = index.cmp uge, %lane, %k24 : index + %n0_k_c6_lt = index.cmp ult, %lane, %k28 : index + %n0_k_c6_ge_v = scf.select %n0_k_c6_ge, %n0_k_local8, %f0 : f32 + %n0_k_c6_in = scf.select %n0_k_c6_lt, %n0_k_c6_ge_v, %f0 : f32 + %n0_k_w6 = kernel.subgroup.reduce %n0_k_c6_in : f32 + %n0_k_c5_ge = index.cmp uge, %lane, %k20 : index + %n0_k_c5_lt = index.cmp ult, %lane, %k24 : index + %n0_k_c5_ge_v = scf.select %n0_k_c5_ge, %n0_k_local8, %f0 : f32 + %n0_k_c5_in = scf.select %n0_k_c5_lt, %n0_k_c5_ge_v, %f0 : f32 + %n0_k_w5 = kernel.subgroup.reduce %n0_k_c5_in : f32 + %n0_k_c7_ge = index.cmp uge, %lane, %k28 : index + %n0_k_c7_lt = index.cmp ult, %lane, %k32 : index + %n0_k_c7_ge_v = scf.select %n0_k_c7_ge, %n0_k_local8, %f0 : f32 + %n0_k_c7_in = scf.select %n0_k_c7_lt, %n0_k_c7_ge_v, %f0 : f32 + %n0_k_w7 = kernel.subgroup.reduce %n0_k_c7_in : f32 + %n0_k_lo02 = scalar.addf %n0_k_w0, %n0_k_w2 : f32 + %n0_k_lo13 = scalar.addf %n0_k_w1, %n0_k_w3 : f32 + %n0_k_lo = scalar.addf %n0_k_lo02, %n0_k_lo13 : f32 + %n0_k_hi46 = scalar.addf %n0_k_w4, %n0_k_w6 : f32 + %n0_k_hi57 = scalar.addf %n0_k_w5, %n0_k_w7 : f32 + %n0_k_hi = scalar.addf %n0_k_hi46, %n0_k_hi57 : f32 + %n0_k_sum = scf.select %lane_lo, %n0_k_lo, %n0_k_hi : f32 + %n0_k_clamped = scalar.maxnumf %n0_k_sum, %norm_eps_sq : f32 + %n0_k_scale = scalar.rsqrtf %n0_k_clamped : f32 + %n0_k_scale_v = vector.splat %n0_k_scale : vector<8xf32> + %n0_kn = vector.mulf %n0_k_scale_v, %kvs0 : vector<8xf32> + %kh0 = vector.fptrunc %n0_kn : vector<8xf32> to vector<8xf16> + %sll0_m = index.mul %stokl, %k136 : index + %sll0_a = index.add %sll0_m, %srow0 : index + %sll0_b = index.add %sll0_a, %k0 : index + %sll0 = index.assume %sll0_b [range(%sll0_b, 0, 2168)] : index + vector.store %kh0, %km_flat[%sll0] : vector<8xf16>, view<2176xf16> + // Introduce Q only after K's payload and partials are dead. + %qvr0 = vector.load %q_view[%sgi0] : view<1073741824xf32> -> vector<8xf32> + %qvs0 = scf.select %stk0_ok, %qvr0, %zv8 : vector<8xf32> + %n0_q_v0 = vector.extract %qvs0[0] : vector<8xf32> -> f32 + %n0_q_sq0 = scalar.mulf %n0_q_v0, %n0_q_v0 : f32 + %n0_q_v1 = vector.extract %qvs0[1] : vector<8xf32> -> f32 + %n0_q_sq1 = scalar.mulf %n0_q_v1, %n0_q_v1 : f32 + %n0_q_v2 = vector.extract %qvs0[2] : vector<8xf32> -> f32 + %n0_q_sq2 = scalar.mulf %n0_q_v2, %n0_q_v2 : f32 + %n0_q_v3 = vector.extract %qvs0[3] : vector<8xf32> -> f32 + %n0_q_sq3 = scalar.mulf %n0_q_v3, %n0_q_v3 : f32 + %n0_q_v4 = vector.extract %qvs0[4] : vector<8xf32> -> f32 + %n0_q_sq4 = scalar.mulf %n0_q_v4, %n0_q_v4 : f32 + %n0_q_v5 = vector.extract %qvs0[5] : vector<8xf32> -> f32 + %n0_q_sq5 = scalar.mulf %n0_q_v5, %n0_q_v5 : f32 + %n0_q_v6 = vector.extract %qvs0[6] : vector<8xf32> -> f32 + %n0_q_sq6 = scalar.mulf %n0_q_v6, %n0_q_v6 : f32 + %n0_q_v7 = vector.extract %qvs0[7] : vector<8xf32> -> f32 + %n0_q_sq7 = scalar.mulf %n0_q_v7, %n0_q_v7 : f32 + %n0_q_p01 = scalar.addf %n0_q_sq0, %n0_q_sq1 : f32 + %n0_q_p23 = scalar.addf %n0_q_sq2, %n0_q_sq3 : f32 + %n0_q_p45 = scalar.addf %n0_q_sq4, %n0_q_sq5 : f32 + %n0_q_p67 = scalar.addf %n0_q_sq6, %n0_q_sq7 : f32 + %n0_q_q03 = scalar.addf %n0_q_p01, %n0_q_p23 : f32 + %n0_q_q47 = scalar.addf %n0_q_p45, %n0_q_p67 : f32 + %n0_q_local8 = scalar.addf %n0_q_q03, %n0_q_q47 : f32 + %n0_q_c0_ge = index.cmp uge, %lane, %k0 : index + %n0_q_c0_lt = index.cmp ult, %lane, %k4 : index + %n0_q_c0_ge_v = scf.select %n0_q_c0_ge, %n0_q_local8, %f0 : f32 + %n0_q_c0_in = scf.select %n0_q_c0_lt, %n0_q_c0_ge_v, %f0 : f32 + %n0_q_w0 = kernel.subgroup.reduce %n0_q_c0_in : f32 + %n0_q_c2_ge = index.cmp uge, %lane, %k8 : index + %n0_q_c2_lt = index.cmp ult, %lane, %k12 : index + %n0_q_c2_ge_v = scf.select %n0_q_c2_ge, %n0_q_local8, %f0 : f32 + %n0_q_c2_in = scf.select %n0_q_c2_lt, %n0_q_c2_ge_v, %f0 : f32 + %n0_q_w2 = kernel.subgroup.reduce %n0_q_c2_in : f32 + %n0_q_c1_ge = index.cmp uge, %lane, %k4 : index + %n0_q_c1_lt = index.cmp ult, %lane, %k8 : index + %n0_q_c1_ge_v = scf.select %n0_q_c1_ge, %n0_q_local8, %f0 : f32 + %n0_q_c1_in = scf.select %n0_q_c1_lt, %n0_q_c1_ge_v, %f0 : f32 + %n0_q_w1 = kernel.subgroup.reduce %n0_q_c1_in : f32 + %n0_q_c3_ge = index.cmp uge, %lane, %k12 : index + %n0_q_c3_lt = index.cmp ult, %lane, %k16 : index + %n0_q_c3_ge_v = scf.select %n0_q_c3_ge, %n0_q_local8, %f0 : f32 + %n0_q_c3_in = scf.select %n0_q_c3_lt, %n0_q_c3_ge_v, %f0 : f32 + %n0_q_w3 = kernel.subgroup.reduce %n0_q_c3_in : f32 + %n0_q_c4_ge = index.cmp uge, %lane, %k16 : index + %n0_q_c4_lt = index.cmp ult, %lane, %k20 : index + %n0_q_c4_ge_v = scf.select %n0_q_c4_ge, %n0_q_local8, %f0 : f32 + %n0_q_c4_in = scf.select %n0_q_c4_lt, %n0_q_c4_ge_v, %f0 : f32 + %n0_q_w4 = kernel.subgroup.reduce %n0_q_c4_in : f32 + %n0_q_c6_ge = index.cmp uge, %lane, %k24 : index + %n0_q_c6_lt = index.cmp ult, %lane, %k28 : index + %n0_q_c6_ge_v = scf.select %n0_q_c6_ge, %n0_q_local8, %f0 : f32 + %n0_q_c6_in = scf.select %n0_q_c6_lt, %n0_q_c6_ge_v, %f0 : f32 + %n0_q_w6 = kernel.subgroup.reduce %n0_q_c6_in : f32 + %n0_q_c5_ge = index.cmp uge, %lane, %k20 : index + %n0_q_c5_lt = index.cmp ult, %lane, %k24 : index + %n0_q_c5_ge_v = scf.select %n0_q_c5_ge, %n0_q_local8, %f0 : f32 + %n0_q_c5_in = scf.select %n0_q_c5_lt, %n0_q_c5_ge_v, %f0 : f32 + %n0_q_w5 = kernel.subgroup.reduce %n0_q_c5_in : f32 + %n0_q_c7_ge = index.cmp uge, %lane, %k28 : index + %n0_q_c7_lt = index.cmp ult, %lane, %k32 : index + %n0_q_c7_ge_v = scf.select %n0_q_c7_ge, %n0_q_local8, %f0 : f32 + %n0_q_c7_in = scf.select %n0_q_c7_lt, %n0_q_c7_ge_v, %f0 : f32 + %n0_q_w7 = kernel.subgroup.reduce %n0_q_c7_in : f32 + %n0_q_lo02 = scalar.addf %n0_q_w0, %n0_q_w2 : f32 + %n0_q_lo13 = scalar.addf %n0_q_w1, %n0_q_w3 : f32 + %n0_q_lo = scalar.addf %n0_q_lo02, %n0_q_lo13 : f32 + %n0_q_hi46 = scalar.addf %n0_q_w4, %n0_q_w6 : f32 + %n0_q_hi57 = scalar.addf %n0_q_w5, %n0_q_w7 : f32 + %n0_q_hi = scalar.addf %n0_q_hi46, %n0_q_hi57 : f32 + %n0_q_sum = scf.select %lane_lo, %n0_q_lo, %n0_q_hi : f32 + %n0_q_clamped = scalar.maxnumf %n0_q_sum, %norm_eps_sq : f32 + %n0_q_scale = scalar.rsqrtf %n0_q_clamped : f32 + %n0_q_scale_v = vector.splat %n0_q_scale : vector<8xf32> + %n0_qn = vector.mulf %n0_q_scale_v, %qvs0 : vector<8xf32> + %qh0 = vector.fptrunc %n0_qn : vector<8xf32> to vector<8xf16> + vector.store %qh0, %qm_flat[%sll0] : vector<8xf16>, view<2176xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + + // K' = K^T for the state update. Retain each row's two f16x8 halves + // in registers until the attention output has vacated Q_main. + %kt_live = index.cmp ult, %tid, %k128 : index + %ktv0, %ktv1 = scf.if %kt_live -> (vector<8xf16>, vector<8xf16>) { + %ktrow = index.assume %tid [range(%tid, 0, 127)] : index + %kts0_a = index.add %ktrow, %k0 : index + %kts0 = index.assume %kts0_a [range(%kts0_a, 0, 2175)] : index + %kt0 = view.load %km_flat[%kts0] : view<2176xf16> -> f16 + %kts1_a = index.add %ktrow, %k136 : index + %kts1 = index.assume %kts1_a [range(%kts1_a, 0, 2175)] : index + %kt1 = view.load %km_flat[%kts1] : view<2176xf16> -> f16 + %kts2_a = index.add %ktrow, %k272 : index + %kts2 = index.assume %kts2_a [range(%kts2_a, 0, 2175)] : index + %kt2 = view.load %km_flat[%kts2] : view<2176xf16> -> f16 + %kts3_a = index.add %ktrow, %k408 : index + %kts3 = index.assume %kts3_a [range(%kts3_a, 0, 2175)] : index + %kt3 = view.load %km_flat[%kts3] : view<2176xf16> -> f16 + %kts4_a = index.add %ktrow, %k544 : index + %kts4 = index.assume %kts4_a [range(%kts4_a, 0, 2175)] : index + %kt4 = view.load %km_flat[%kts4] : view<2176xf16> -> f16 + %kts5_a = index.add %ktrow, %k680 : index + %kts5 = index.assume %kts5_a [range(%kts5_a, 0, 2175)] : index + %kt5 = view.load %km_flat[%kts5] : view<2176xf16> -> f16 + %kts6_a = index.add %ktrow, %k816 : index + %kts6 = index.assume %kts6_a [range(%kts6_a, 0, 2175)] : index + %kt6 = view.load %km_flat[%kts6] : view<2176xf16> -> f16 + %kts7_a = index.add %ktrow, %k952 : index + %kts7 = index.assume %kts7_a [range(%kts7_a, 0, 2175)] : index + %kt7 = view.load %km_flat[%kts7] : view<2176xf16> -> f16 + %kts8_a = index.add %ktrow, %k1088 : index + %kts8 = index.assume %kts8_a [range(%kts8_a, 0, 2175)] : index + %kt8 = view.load %km_flat[%kts8] : view<2176xf16> -> f16 + %kts9_a = index.add %ktrow, %k1224 : index + %kts9 = index.assume %kts9_a [range(%kts9_a, 0, 2175)] : index + %kt9 = view.load %km_flat[%kts9] : view<2176xf16> -> f16 + %kts10_a = index.add %ktrow, %k1360 : index + %kts10 = index.assume %kts10_a [range(%kts10_a, 0, 2175)] : index + %kt10 = view.load %km_flat[%kts10] : view<2176xf16> -> f16 + %kts11_a = index.add %ktrow, %k1496 : index + %kts11 = index.assume %kts11_a [range(%kts11_a, 0, 2175)] : index + %kt11 = view.load %km_flat[%kts11] : view<2176xf16> -> f16 + %kts12_a = index.add %ktrow, %k1632 : index + %kts12 = index.assume %kts12_a [range(%kts12_a, 0, 2175)] : index + %kt12 = view.load %km_flat[%kts12] : view<2176xf16> -> f16 + %kts13_a = index.add %ktrow, %k1768 : index + %kts13 = index.assume %kts13_a [range(%kts13_a, 0, 2175)] : index + %kt13 = view.load %km_flat[%kts13] : view<2176xf16> -> f16 + %kts14_a = index.add %ktrow, %k1904 : index + %kts14 = index.assume %kts14_a [range(%kts14_a, 0, 2175)] : index + %kt14 = view.load %km_flat[%kts14] : view<2176xf16> -> f16 + %kts15_a = index.add %ktrow, %k2040 : index + %kts15 = index.assume %kts15_a [range(%kts15_a, 0, 2175)] : index + %kt15 = view.load %km_flat[%kts15] : view<2176xf16> -> f16 + %ktlo = vector.from_elements %kt0, %kt1, %kt2, %kt3, %kt4, %kt5, %kt6, %kt7 : vector<8xf16> + %kthi = vector.from_elements %kt8, %kt9, %kt10, %kt11, %kt12, %kt13, %kt14, %kt15 : vector<8xf16> + scf.yield %ktlo, %kthi : vector<8xf16>, vector<8xf16> + } else { + scf.yield %fzh8, %fzh8 : vector<8xf16>, vector<8xf16> + } + + // State shard -> LDS as f16: the lhs of both A and B. + %shx0 = vector.fptrunc %sa0 : vector<4xf32> to vector<4xf16> + %shi0_m = index.mul %k0, %k136 : index + %shi0_a = index.add %shi0_m, %row0 : index + %shi0 = index.assume %shi0_a [range(%shi0_a, 0, 2172)] : index + vector.store %shx0, %sh_flat[%shi0] : vector<4xf16>, view<2176xf16> + %shx1 = vector.fptrunc %sa1 : vector<4xf32> to vector<4xf16> + %shi1_m = index.mul %k1, %k136 : index + %shi1_a = index.add %shi1_m, %row0 : index + %shi1 = index.assume %shi1_a [range(%shi1_a, 0, 2172)] : index + vector.store %shx1, %sh_flat[%shi1] : vector<4xf16>, view<2176xf16> + %shx2 = vector.fptrunc %sa2 : vector<4xf32> to vector<4xf16> + %shi2_m = index.mul %k2, %k136 : index + %shi2_a = index.add %shi2_m, %row0 : index + %shi2 = index.assume %shi2_a [range(%shi2_a, 0, 2172)] : index + vector.store %shx2, %sh_flat[%shi2] : vector<4xf16>, view<2176xf16> + %shx3 = vector.fptrunc %sa3 : vector<4xf32> to vector<4xf16> + %shi3_m = index.mul %k3, %k136 : index + %shi3_a = index.add %shi3_m, %row0 : index + %shi3 = index.assume %shi3_a [range(%shi3_a, 0, 2172)] : index + vector.store %shx3, %sh_flat[%shi3] : vector<4xf16>, view<2176xf16> + %shx4 = vector.fptrunc %sa4 : vector<4xf32> to vector<4xf16> + %shi4_m = index.mul %k4, %k136 : index + %shi4_a = index.add %shi4_m, %row0 : index + %shi4 = index.assume %shi4_a [range(%shi4_a, 0, 2172)] : index + vector.store %shx4, %sh_flat[%shi4] : vector<4xf16>, view<2176xf16> + %shx5 = vector.fptrunc %sa5 : vector<4xf32> to vector<4xf16> + %shi5_m = index.mul %k5, %k136 : index + %shi5_a = index.add %shi5_m, %row0 : index + %shi5 = index.assume %shi5_a [range(%shi5_a, 0, 2172)] : index + vector.store %shx5, %sh_flat[%shi5] : vector<4xf16>, view<2176xf16> + %shx6 = vector.fptrunc %sa6 : vector<4xf32> to vector<4xf16> + %shi6_m = index.mul %k6, %k136 : index + %shi6_a = index.add %shi6_m, %row0 : index + %shi6 = index.assume %shi6_a [range(%shi6_a, 0, 2172)] : index + vector.store %shx6, %sh_flat[%shi6] : vector<4xf16>, view<2176xf16> + %shx7 = vector.fptrunc %sa7 : vector<4xf32> to vector<4xf16> + %shi7_m = index.mul %k7, %k136 : index + %shi7_a = index.add %shi7_m, %row0 : index + %shi7 = index.assume %shi7_a [range(%shi7_a, 0, 2172)] : index + vector.store %shx7, %sh_flat[%shi7] : vector<4xf16>, view<2176xf16> + %shx8 = vector.fptrunc %sa8 : vector<4xf32> to vector<4xf16> + %shi8_m = index.mul %k8, %k136 : index + %shi8_a = index.add %shi8_m, %row0 : index + %shi8 = index.assume %shi8_a [range(%shi8_a, 0, 2172)] : index + vector.store %shx8, %sh_flat[%shi8] : vector<4xf16>, view<2176xf16> + %shx9 = vector.fptrunc %sa9 : vector<4xf32> to vector<4xf16> + %shi9_m = index.mul %k9, %k136 : index + %shi9_a = index.add %shi9_m, %row0 : index + %shi9 = index.assume %shi9_a [range(%shi9_a, 0, 2172)] : index + vector.store %shx9, %sh_flat[%shi9] : vector<4xf16>, view<2176xf16> + %shx10 = vector.fptrunc %sa10 : vector<4xf32> to vector<4xf16> + %shi10_m = index.mul %k10, %k136 : index + %shi10_a = index.add %shi10_m, %row0 : index + %shi10 = index.assume %shi10_a [range(%shi10_a, 0, 2172)] : index + vector.store %shx10, %sh_flat[%shi10] : vector<4xf16>, view<2176xf16> + %shx11 = vector.fptrunc %sa11 : vector<4xf32> to vector<4xf16> + %shi11_m = index.mul %k11, %k136 : index + %shi11_a = index.add %shi11_m, %row0 : index + %shi11 = index.assume %shi11_a [range(%shi11_a, 0, 2172)] : index + vector.store %shx11, %sh_flat[%shi11] : vector<4xf16>, view<2176xf16> + %shx12 = vector.fptrunc %sa12 : vector<4xf32> to vector<4xf16> + %shi12_m = index.mul %k12, %k136 : index + %shi12_a = index.add %shi12_m, %row0 : index + %shi12 = index.assume %shi12_a [range(%shi12_a, 0, 2172)] : index + vector.store %shx12, %sh_flat[%shi12] : vector<4xf16>, view<2176xf16> + %shx13 = vector.fptrunc %sa13 : vector<4xf32> to vector<4xf16> + %shi13_m = index.mul %k13, %k136 : index + %shi13_a = index.add %shi13_m, %row0 : index + %shi13 = index.assume %shi13_a [range(%shi13_a, 0, 2172)] : index + vector.store %shx13, %sh_flat[%shi13] : vector<4xf16>, view<2176xf16> + %shx14 = vector.fptrunc %sa14 : vector<4xf32> to vector<4xf16> + %shi14_m = index.mul %k14, %k136 : index + %shi14_a = index.add %shi14_m, %row0 : index + %shi14 = index.assume %shi14_a [range(%shi14_a, 0, 2172)] : index + vector.store %shx14, %sh_flat[%shi14] : vector<4xf16>, view<2176xf16> + %shx15 = vector.fptrunc %sa15 : vector<4xf32> to vector<4xf16> + %shi15_m = index.mul %k15, %k136 : index + %shi15_a = index.add %shi15_m, %row0 : index + %shi15 = index.assume %shi15_a [range(%shi15_a, 0, 2172)] : index + vector.store %shx15, %sh_flat[%shi15] : vector<4xf16>, view<2176xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + + // ---- matmuls ---------------------------------------------------------- + %ma_a0 = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %mb_a0 = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %ma_l0 = vector.fragment.load %sh_lhs[%k0, %k0] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r0 = vector.fragment.load %km_rhs[%k0, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a1 = vector.mma %ma_l0, %ma_r0, %ma_a0 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r0 = vector.fragment.load %qm_rhs[%k0, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a1 = vector.mma %ma_l0, %mb_r0, %mb_a0 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %ma_l1 = vector.fragment.load %sh_lhs[%k0, %k16] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r1 = vector.fragment.load %km_rhs[%k16, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a2 = vector.mma %ma_l1, %ma_r1, %ma_a1 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r1 = vector.fragment.load %qm_rhs[%k16, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a2 = vector.mma %ma_l1, %mb_r1, %mb_a1 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %ma_l2 = vector.fragment.load %sh_lhs[%k0, %k32] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r2 = vector.fragment.load %km_rhs[%k32, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a3 = vector.mma %ma_l2, %ma_r2, %ma_a2 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r2 = vector.fragment.load %qm_rhs[%k32, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a3 = vector.mma %ma_l2, %mb_r2, %mb_a2 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %ma_l3 = vector.fragment.load %sh_lhs[%k0, %k48] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r3 = vector.fragment.load %km_rhs[%k48, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a4 = vector.mma %ma_l3, %ma_r3, %ma_a3 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r3 = vector.fragment.load %qm_rhs[%k48, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a4 = vector.mma %ma_l3, %mb_r3, %mb_a3 : vector<16xf16>, vector<16xf16>, vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + %ma_l4 = vector.fragment.load %sh_lhs[%k0, %k64] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r4 = vector.fragment.load %km_rhs[%k64, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a5 = vector.mma %ma_l4, %ma_r4, %ma_a4 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r4 = vector.fragment.load %qm_rhs[%k64, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a5 = vector.mma %ma_l4, %mb_r4, %mb_a4 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %ma_l5 = vector.fragment.load %sh_lhs[%k0, %k80] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r5 = vector.fragment.load %km_rhs[%k80, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a6 = vector.mma %ma_l5, %ma_r5, %ma_a5 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r5 = vector.fragment.load %qm_rhs[%k80, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a6 = vector.mma %ma_l5, %mb_r5, %mb_a5 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %ma_l6 = vector.fragment.load %sh_lhs[%k0, %k96] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r6 = vector.fragment.load %km_rhs[%k96, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a7 = vector.mma %ma_l6, %ma_r6, %ma_a6 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r6 = vector.fragment.load %qm_rhs[%k96, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a7 = vector.mma %ma_l6, %mb_r6, %mb_a6 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %ma_l7 = vector.fragment.load %sh_lhs[%k0, %k112] shape [%k16, %k16] : view<16x128xf16, %lay_colrow> -> vector<16xf16> + %ma_r7 = vector.fragment.load %km_rhs[%k112, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %ma_a8 = vector.mma %ma_l7, %ma_r7, %ma_a7 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mb_r7 = vector.fragment.load %qm_rhs[%k112, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mb_a8 = vector.mma %ma_l7, %mb_r7, %mb_a7 : vector<16xf16>, vector<16xf16>, vector<8xf32> + vector.fragment.store %ma_a8, %a_res[%k0, %k0] shape [%k16, %k16] : vector<8xf32>, view<16x16xf32, %lay_tile> + vector.fragment.store %mb_a8, %b_res[%k0, %k0] shape [%k16, %k16] : vector<8xf32>, view<16x16xf32, %lay_tile> + // G and QK are per head, so the two waves take one each. + scf.if %is_w0 { + %mg_a0 = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %mg_l0 = vector.fragment.load %km_lhs[%k0, %k0] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r0 = vector.fragment.load %km_rhs[%k0, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a1 = vector.mma %mg_l0, %mg_r0, %mg_a0 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l1 = vector.fragment.load %km_lhs[%k0, %k16] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r1 = vector.fragment.load %km_rhs[%k16, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a2 = vector.mma %mg_l1, %mg_r1, %mg_a1 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l2 = vector.fragment.load %km_lhs[%k0, %k32] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r2 = vector.fragment.load %km_rhs[%k32, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a3 = vector.mma %mg_l2, %mg_r2, %mg_a2 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l3 = vector.fragment.load %km_lhs[%k0, %k48] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r3 = vector.fragment.load %km_rhs[%k48, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a4 = vector.mma %mg_l3, %mg_r3, %mg_a3 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l4 = vector.fragment.load %km_lhs[%k0, %k64] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r4 = vector.fragment.load %km_rhs[%k64, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a5 = vector.mma %mg_l4, %mg_r4, %mg_a4 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l5 = vector.fragment.load %km_lhs[%k0, %k80] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r5 = vector.fragment.load %km_rhs[%k80, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a6 = vector.mma %mg_l5, %mg_r5, %mg_a5 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l6 = vector.fragment.load %km_lhs[%k0, %k96] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r6 = vector.fragment.load %km_rhs[%k96, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a7 = vector.mma %mg_l6, %mg_r6, %mg_a6 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mg_l7 = vector.fragment.load %km_lhs[%k0, %k112] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mg_r7 = vector.fragment.load %km_rhs[%k112, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mg_a8 = vector.mma %mg_l7, %mg_r7, %mg_a7 : vector<16xf16>, vector<16xf16>, vector<8xf32> + vector.fragment.store %mg_a8, %gr_res[%k0, %k0] shape [%k16, %k16] : vector<8xf32>, view<16x16xf32, %lay_gram> + } + scf.if %is_w1 { + %mq_a0 = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %mq_l0 = vector.fragment.load %km_lhs[%k0, %k0] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r0 = vector.fragment.load %qm_rhs[%k0, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a1 = vector.mma %mq_l0, %mq_r0, %mq_a0 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l1 = vector.fragment.load %km_lhs[%k0, %k16] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r1 = vector.fragment.load %qm_rhs[%k16, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a2 = vector.mma %mq_l1, %mq_r1, %mq_a1 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l2 = vector.fragment.load %km_lhs[%k0, %k32] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r2 = vector.fragment.load %qm_rhs[%k32, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a3 = vector.mma %mq_l2, %mq_r2, %mq_a2 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l3 = vector.fragment.load %km_lhs[%k0, %k48] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r3 = vector.fragment.load %qm_rhs[%k48, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a4 = vector.mma %mq_l3, %mq_r3, %mq_a3 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l4 = vector.fragment.load %km_lhs[%k0, %k64] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r4 = vector.fragment.load %qm_rhs[%k64, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a5 = vector.mma %mq_l4, %mq_r4, %mq_a4 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l5 = vector.fragment.load %km_lhs[%k0, %k80] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r5 = vector.fragment.load %qm_rhs[%k80, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a6 = vector.mma %mq_l5, %mq_r5, %mq_a5 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l6 = vector.fragment.load %km_lhs[%k0, %k96] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r6 = vector.fragment.load %qm_rhs[%k96, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a7 = vector.mma %mq_l6, %mq_r6, %mq_a6 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %mq_l7 = vector.fragment.load %km_lhs[%k0, %k112] shape [%k16, %k16] : view<16x128xf16, %lay_tokrow> -> vector<16xf16> + %mq_r7 = vector.fragment.load %qm_rhs[%k112, %k0] shape [%k16, %k16] : view<128x16xf16, %lay_rowtok> -> vector<16xf16> + %mq_a8 = vector.mma %mq_l7, %mq_r7, %mq_a7 : vector<16xf16>, vector<16xf16>, vector<8xf32> + vector.fragment.store %mq_a8, %qkt_res[%k0, %k0] shape [%k16, %k16] : vector<8xf32>, view<16x16xf32, %lay_gram> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %df0_a = index.add %tid, %k0 : index + %df0 = index.assume %df0_a [range(%df0_a, 0, 255)] : index + %ds0_r = index.div %df0, %k16 : index + %ds0 = index.assume %ds0_r [range(%ds0_r, 0, 15)] : index + %dt0_r = index.rem %df0, %k16 : index + %dt0 = index.assume %dt0_r [range(%dt0_r, 0, 15)] : index + %du0 = index.cmp ult, %ds0, %dt0 : index + scf.if %du0 { + %dsi0_m = index.mul %ds0, %k4 : index + %dsi0_a = index.add %dsi0_m, %k3 : index + %dsi0 = index.assume %dsi0_a [range(%dsi0_a, 3, 63)] : index + %dti0_m = index.mul %dt0, %k4 : index + %dti0_a = index.add %dti0_m, %k3 : index + %dti0 = index.assume %dti0_a [range(%dti0_a, 3, 63)] : index + %dsl0 = view.load %gs_flat[%dsi0] : view<80xf32> -> f32 + %dtl0 = view.load %gs_flat[%dti0] : view<80xf32> -> f32 + %ddl0 = scalar.subf %dtl0, %dsl0 : f32 + %ddc0 = scalar.expf %ddl0 : f32 + %dgv0 = view.load %gr_flat[%df0] : view<256xf32> -> f32 + %dqv0 = view.load %qkt_flat[%df0] : view<256xf32> -> f32 + %dgs0 = scalar.mulf %dgv0, %ddc0 : f32 + %dqs0 = scalar.mulf %dqv0, %ddc0 : f32 + view.store %dgs0, %gr_flat[%df0] : f32, view<256xf32> + view.store %dqs0, %qkt_flat[%df0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %vlt0_r = index.add %t0, %k0 : index + %vlt0_s = index.add %vlt0_r, %vs_tok : index + %vlt0_ok = index.cmp ult, %vlt0_s, %n_tokens : index + %vlt0_c = scf.select %vlt0_ok, %vlt0_s, %c0 : index + %vlt0 = index.assume %vlt0_c [range(%vlt0_c, 0, 1048575)] : index + %vli0_m = index.mul %vlt0, %sv2 : index + %vli0_a = index.add %v_head_base, %vli0_m : index + %vli0_c = index.add %vli0_a, %vs_gcol : index + %vli0 = index.assume %vli0_c [range(%vli0_c, 0, 1073741823)] : index + %vlv0 = view.load %v_view[%vli0] : view<1073741824xf32> -> f32 + %vls0_m = index.mul %vs_tok, %k128 : index + %vls0_a = index.add %vls0_m, %vs_col : index + %vls0_b = index.add %vls0_a, %k0 : index + %vls0 = index.assume %vls0_b [range(%vls0_b, 0, 2047)] : index + view.store %vlv0, %vst_flat[%vls0] : f32, view<2048xf32> + %vlt1_r = index.add %t0, %k2 : index + %vlt1_s = index.add %vlt1_r, %vs_tok : index + %vlt1_ok = index.cmp ult, %vlt1_s, %n_tokens : index + %vlt1_c = scf.select %vlt1_ok, %vlt1_s, %c0 : index + %vlt1 = index.assume %vlt1_c [range(%vlt1_c, 0, 1048575)] : index + %vli1_m = index.mul %vlt1, %sv2 : index + %vli1_a = index.add %v_head_base, %vli1_m : index + %vli1_c = index.add %vli1_a, %vs_gcol : index + %vli1 = index.assume %vli1_c [range(%vli1_c, 0, 1073741823)] : index + %vlv1 = view.load %v_view[%vli1] : view<1073741824xf32> -> f32 + %vls1_m = index.mul %vs_tok, %k128 : index + %vls1_a = index.add %vls1_m, %vs_col : index + %vls1_b = index.add %vls1_a, %k256 : index + %vls1 = index.assume %vls1_b [range(%vls1_b, 0, 2047)] : index + view.store %vlv1, %vst_flat[%vls1] : f32, view<2048xf32> + %vlt2_r = index.add %t0, %k4 : index + %vlt2_s = index.add %vlt2_r, %vs_tok : index + %vlt2_ok = index.cmp ult, %vlt2_s, %n_tokens : index + %vlt2_c = scf.select %vlt2_ok, %vlt2_s, %c0 : index + %vlt2 = index.assume %vlt2_c [range(%vlt2_c, 0, 1048575)] : index + %vli2_m = index.mul %vlt2, %sv2 : index + %vli2_a = index.add %v_head_base, %vli2_m : index + %vli2_c = index.add %vli2_a, %vs_gcol : index + %vli2 = index.assume %vli2_c [range(%vli2_c, 0, 1073741823)] : index + %vlv2 = view.load %v_view[%vli2] : view<1073741824xf32> -> f32 + %vls2_m = index.mul %vs_tok, %k128 : index + %vls2_a = index.add %vls2_m, %vs_col : index + %vls2_b = index.add %vls2_a, %k512 : index + %vls2 = index.assume %vls2_b [range(%vls2_b, 0, 2047)] : index + view.store %vlv2, %vst_flat[%vls2] : f32, view<2048xf32> + %vlt3_r = index.add %t0, %k6 : index + %vlt3_s = index.add %vlt3_r, %vs_tok : index + %vlt3_ok = index.cmp ult, %vlt3_s, %n_tokens : index + %vlt3_c = scf.select %vlt3_ok, %vlt3_s, %c0 : index + %vlt3 = index.assume %vlt3_c [range(%vlt3_c, 0, 1048575)] : index + %vli3_m = index.mul %vlt3, %sv2 : index + %vli3_a = index.add %v_head_base, %vli3_m : index + %vli3_c = index.add %vli3_a, %vs_gcol : index + %vli3 = index.assume %vli3_c [range(%vli3_c, 0, 1073741823)] : index + %vlv3 = view.load %v_view[%vli3] : view<1073741824xf32> -> f32 + %vls3_m = index.mul %vs_tok, %k128 : index + %vls3_a = index.add %vls3_m, %vs_col : index + %vls3_b = index.add %vls3_a, %k768 : index + %vls3 = index.assume %vls3_b [range(%vls3_b, 0, 2047)] : index + view.store %vlv3, %vst_flat[%vls3] : f32, view<2048xf32> + %vlt4_r = index.add %t0, %k8 : index + %vlt4_s = index.add %vlt4_r, %vs_tok : index + %vlt4_ok = index.cmp ult, %vlt4_s, %n_tokens : index + %vlt4_c = scf.select %vlt4_ok, %vlt4_s, %c0 : index + %vlt4 = index.assume %vlt4_c [range(%vlt4_c, 0, 1048575)] : index + %vli4_m = index.mul %vlt4, %sv2 : index + %vli4_a = index.add %v_head_base, %vli4_m : index + %vli4_c = index.add %vli4_a, %vs_gcol : index + %vli4 = index.assume %vli4_c [range(%vli4_c, 0, 1073741823)] : index + %vlv4 = view.load %v_view[%vli4] : view<1073741824xf32> -> f32 + %vls4_m = index.mul %vs_tok, %k128 : index + %vls4_a = index.add %vls4_m, %vs_col : index + %vls4_b = index.add %vls4_a, %k1024 : index + %vls4 = index.assume %vls4_b [range(%vls4_b, 0, 2047)] : index + view.store %vlv4, %vst_flat[%vls4] : f32, view<2048xf32> + %vlt5_r = index.add %t0, %k10 : index + %vlt5_s = index.add %vlt5_r, %vs_tok : index + %vlt5_ok = index.cmp ult, %vlt5_s, %n_tokens : index + %vlt5_c = scf.select %vlt5_ok, %vlt5_s, %c0 : index + %vlt5 = index.assume %vlt5_c [range(%vlt5_c, 0, 1048575)] : index + %vli5_m = index.mul %vlt5, %sv2 : index + %vli5_a = index.add %v_head_base, %vli5_m : index + %vli5_c = index.add %vli5_a, %vs_gcol : index + %vli5 = index.assume %vli5_c [range(%vli5_c, 0, 1073741823)] : index + %vlv5 = view.load %v_view[%vli5] : view<1073741824xf32> -> f32 + %vls5_m = index.mul %vs_tok, %k128 : index + %vls5_a = index.add %vls5_m, %vs_col : index + %vls5_b = index.add %vls5_a, %k1280 : index + %vls5 = index.assume %vls5_b [range(%vls5_b, 0, 2047)] : index + view.store %vlv5, %vst_flat[%vls5] : f32, view<2048xf32> + %vlt6_r = index.add %t0, %k12 : index + %vlt6_s = index.add %vlt6_r, %vs_tok : index + %vlt6_ok = index.cmp ult, %vlt6_s, %n_tokens : index + %vlt6_c = scf.select %vlt6_ok, %vlt6_s, %c0 : index + %vlt6 = index.assume %vlt6_c [range(%vlt6_c, 0, 1048575)] : index + %vli6_m = index.mul %vlt6, %sv2 : index + %vli6_a = index.add %v_head_base, %vli6_m : index + %vli6_c = index.add %vli6_a, %vs_gcol : index + %vli6 = index.assume %vli6_c [range(%vli6_c, 0, 1073741823)] : index + %vlv6 = view.load %v_view[%vli6] : view<1073741824xf32> -> f32 + %vls6_m = index.mul %vs_tok, %k128 : index + %vls6_a = index.add %vls6_m, %vs_col : index + %vls6_b = index.add %vls6_a, %k1536 : index + %vls6 = index.assume %vls6_b [range(%vls6_b, 0, 2047)] : index + view.store %vlv6, %vst_flat[%vls6] : f32, view<2048xf32> + %vlt7_r = index.add %t0, %k14 : index + %vlt7_s = index.add %vlt7_r, %vs_tok : index + %vlt7_ok = index.cmp ult, %vlt7_s, %n_tokens : index + %vlt7_c = scf.select %vlt7_ok, %vlt7_s, %c0 : index + %vlt7 = index.assume %vlt7_c [range(%vlt7_c, 0, 1048575)] : index + %vli7_m = index.mul %vlt7, %sv2 : index + %vli7_a = index.add %v_head_base, %vli7_m : index + %vli7_c = index.add %vli7_a, %vs_gcol : index + %vli7 = index.assume %vli7_c [range(%vli7_c, 0, 1073741823)] : index + %vlv7 = view.load %v_view[%vli7] : view<1073741824xf32> -> f32 + %vls7_m = index.mul %vs_tok, %k128 : index + %vls7_a = index.add %vls7_m, %vs_col : index + %vls7_b = index.add %vls7_a, %k1792 : index + %vls7 = index.assume %vls7_b [range(%vls7_b, 0, 2047)] : index + view.store %vlv7, %vst_flat[%vls7] : f32, view<2048xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + + // Lane c owns column c; recurrence operands are registers or LDS broadcasts with no cross-lane exchange. + %gq0 = vector.load %gs_flat[%k0] : view<80xf32> -> vector<4xf32> + %bv0 = vector.extract %gq0[0] : vector<4xf32> -> f32 + %cp0 = vector.extract %gq0[1] : vector<4xf32> -> f32 + %ev0 = vector.extract %gq0[2] : vector<4xf32> -> f32 + %tk0_r = index.add %t0, %k0 : index + %tk0_ok = index.cmp ult, %tk0_r, %n_tokens : index + %tk0_s = scf.select %tk0_ok, %tk0_r, %c0 : index + %tk0 = index.assume %tk0_s [range(%tk0_s, 0, 1048575)] : index + %av0_i = index.add %tile_row, %k0 : index + %av0_b = index.assume %av0_i [range(%av0_i, 0, 271)] : index + %av0 = view.load %a_flat[%av0_b] : view<272xf32> -> f32 + %bvv0 = view.load %b_flat[%av0_b] : view<272xf32> -> f32 + %vrd0_a = index.add %rec_col, %k0 : index + %vrd0 = index.assume %vrd0_a [range(%vrd0_a, 0, 2047)] : index + %vval0 = view.load %vst_flat[%vrd0] : view<2048xf32> -> f32 + %adec0 = scalar.mulf %av0, %cp0 : f32 + %vsub0 = scalar.subf %vval0, %adec0 : f32 + // Keep the original padded multiply, including its signed zero. + %active_index0 = index.add %t0, %k0 : index + %row_in_range0 = index.cmp ult, %active_index0, %n_tokens : index + %active0 = scalar.ori %row_in_range0, %all_recurrence_rows : i1 + %dl0 = scf.if %active0 -> (f32) { + %value = scalar.mulf %vsub0, %bv0 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval0, %bv0 : f32 + scf.yield %padded : f32 + } + %dh0 = scalar.mulf %dl0, %ev0 : f32 + %grw0_0_a = vector.load %gr_flat[%k0] : view<256xf32> -> vector<4xf32> + %grw0_0_bi = index.constant 4 : index + %grw0_0_b = vector.load %gr_flat[%grw0_0_bi] : view<256xf32> -> vector<4xf32> + %grw0_1_a = vector.load %gr_flat[%k8] : view<256xf32> -> vector<4xf32> + %grw0_1_bi = index.constant 12 : index + %grw0_1_b = vector.load %gr_flat[%grw0_1_bi] : view<256xf32> -> vector<4xf32> + %qkw0_0_a = vector.load %qkt_flat[%k0] : view<256xf32> -> vector<4xf32> + %qkw0_0_bi = index.constant 4 : index + %qkw0_0_b = vector.load %qkt_flat[%qkw0_0_bi] : view<256xf32> -> vector<4xf32> + %qkw0_1_a = vector.load %qkt_flat[%k8] : view<256xf32> -> vector<4xf32> + %qkw0_1_bi = index.constant 12 : index + %qkw0_1_b = vector.load %qkt_flat[%qkw0_1_bi] : view<256xf32> -> vector<4xf32> + %qkd0_0 = vector.extract %qkw0_0_a[0] : vector<4xf32> -> f32 + %qsum0 = scalar.mulf %qkd0_0, %dl0 : f32 + %ao0_r = scalar.fmaf %bvv0, %cp0, %qsum0 : f32 + %ao0 = scf.if %active0 -> (f32) { + %value = scalar.mulf %ao0_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi0_a = index.add %rec_col, %k0 : index + %aoi0 = index.assume %aoi0_a [range(%aoi0_a, 0, 2047)] : index + %ddh0 = scalar.fptrunc %dh0 : f32 to f16 + %publish_active0 = scalar.andi %lane_live, %active0 : i1 + scf.if %publish_active0 { + view.store %ao0, %aost_flat[%aoi0] : f32, view<2048xf32> + } + %gp0_1 = vector.extract %grw0_0_a[1] : vector<4xf32> -> f32 + %ag0_1 = scalar.mulf %gp0_1, %dl0 : f32 + %qp0_1 = vector.extract %qkw0_0_a[1] : vector<4xf32> -> f32 + %aq0_1 = scalar.mulf %qp0_1, %dl0 : f32 + %gp0_2 = vector.extract %grw0_0_a[2] : vector<4xf32> -> f32 + %ag0_2 = scalar.mulf %gp0_2, %dl0 : f32 + %qp0_2 = vector.extract %qkw0_0_a[2] : vector<4xf32> -> f32 + %aq0_2 = scalar.mulf %qp0_2, %dl0 : f32 + %gp0_3 = vector.extract %grw0_0_a[3] : vector<4xf32> -> f32 + %ag0_3 = scalar.mulf %gp0_3, %dl0 : f32 + %qp0_3 = vector.extract %qkw0_0_a[3] : vector<4xf32> -> f32 + %aq0_3 = scalar.mulf %qp0_3, %dl0 : f32 + %gp0_4 = vector.extract %grw0_0_b[0] : vector<4xf32> -> f32 + %ag0_4 = scalar.mulf %gp0_4, %dl0 : f32 + %qp0_4 = vector.extract %qkw0_0_b[0] : vector<4xf32> -> f32 + %aq0_4 = scalar.mulf %qp0_4, %dl0 : f32 + %gp0_5 = vector.extract %grw0_0_b[1] : vector<4xf32> -> f32 + %ag0_5 = scalar.mulf %gp0_5, %dl0 : f32 + %qp0_5 = vector.extract %qkw0_0_b[1] : vector<4xf32> -> f32 + %aq0_5 = scalar.mulf %qp0_5, %dl0 : f32 + %gp0_6 = vector.extract %grw0_0_b[2] : vector<4xf32> -> f32 + %ag0_6 = scalar.mulf %gp0_6, %dl0 : f32 + %qp0_6 = vector.extract %qkw0_0_b[2] : vector<4xf32> -> f32 + %aq0_6 = scalar.mulf %qp0_6, %dl0 : f32 + %gp0_7 = vector.extract %grw0_0_b[3] : vector<4xf32> -> f32 + %ag0_7 = scalar.mulf %gp0_7, %dl0 : f32 + %qp0_7 = vector.extract %qkw0_0_b[3] : vector<4xf32> -> f32 + %aq0_7 = scalar.mulf %qp0_7, %dl0 : f32 + %gp0_8 = vector.extract %grw0_1_a[0] : vector<4xf32> -> f32 + %ag0_8 = scalar.mulf %gp0_8, %dl0 : f32 + %qp0_8 = vector.extract %qkw0_1_a[0] : vector<4xf32> -> f32 + %aq0_8 = scalar.mulf %qp0_8, %dl0 : f32 + %gp0_9 = vector.extract %grw0_1_a[1] : vector<4xf32> -> f32 + %ag0_9 = scalar.mulf %gp0_9, %dl0 : f32 + %qp0_9 = vector.extract %qkw0_1_a[1] : vector<4xf32> -> f32 + %aq0_9 = scalar.mulf %qp0_9, %dl0 : f32 + %gp0_10 = vector.extract %grw0_1_a[2] : vector<4xf32> -> f32 + %ag0_10 = scalar.mulf %gp0_10, %dl0 : f32 + %qp0_10 = vector.extract %qkw0_1_a[2] : vector<4xf32> -> f32 + %aq0_10 = scalar.mulf %qp0_10, %dl0 : f32 + %gp0_11 = vector.extract %grw0_1_a[3] : vector<4xf32> -> f32 + %ag0_11 = scalar.mulf %gp0_11, %dl0 : f32 + %qp0_11 = vector.extract %qkw0_1_a[3] : vector<4xf32> -> f32 + %aq0_11 = scalar.mulf %qp0_11, %dl0 : f32 + %gp0_12 = vector.extract %grw0_1_b[0] : vector<4xf32> -> f32 + %ag0_12 = scalar.mulf %gp0_12, %dl0 : f32 + %qp0_12 = vector.extract %qkw0_1_b[0] : vector<4xf32> -> f32 + %aq0_12 = scalar.mulf %qp0_12, %dl0 : f32 + %gp0_13 = vector.extract %grw0_1_b[1] : vector<4xf32> -> f32 + %ag0_13 = scalar.mulf %gp0_13, %dl0 : f32 + %qp0_13 = vector.extract %qkw0_1_b[1] : vector<4xf32> -> f32 + %aq0_13 = scalar.mulf %qp0_13, %dl0 : f32 + %gp0_14 = vector.extract %grw0_1_b[2] : vector<4xf32> -> f32 + %ag0_14 = scalar.mulf %gp0_14, %dl0 : f32 + %qp0_14 = vector.extract %qkw0_1_b[2] : vector<4xf32> -> f32 + %aq0_14 = scalar.mulf %qp0_14, %dl0 : f32 + %gp0_15 = vector.extract %grw0_1_b[3] : vector<4xf32> -> f32 + %ag0_15 = scalar.mulf %gp0_15, %dl0 : f32 + %qp0_15 = vector.extract %qkw0_1_b[3] : vector<4xf32> -> f32 + %aq0_15 = scalar.mulf %qp0_15, %dl0 : f32 + %gq1 = vector.load %gs_flat[%k4] : view<80xf32> -> vector<4xf32> + %bv1 = vector.extract %gq1[0] : vector<4xf32> -> f32 + %cp1 = vector.extract %gq1[1] : vector<4xf32> -> f32 + %ev1 = vector.extract %gq1[2] : vector<4xf32> -> f32 + %tk1_r = index.add %t0, %k1 : index + %tk1_ok = index.cmp ult, %tk1_r, %n_tokens : index + %tk1_s = scf.select %tk1_ok, %tk1_r, %c0 : index + %tk1 = index.assume %tk1_s [range(%tk1_s, 0, 1048575)] : index + %av1_i = index.add %tile_row, %k1 : index + %av1_b = index.assume %av1_i [range(%av1_i, 0, 271)] : index + %av1 = view.load %a_flat[%av1_b] : view<272xf32> -> f32 + %bvv1 = view.load %b_flat[%av1_b] : view<272xf32> -> f32 + %vrd1_a = index.add %rec_col, %k128 : index + %vrd1 = index.assume %vrd1_a [range(%vrd1_a, 0, 2047)] : index + %vval1 = view.load %vst_flat[%vrd1] : view<2048xf32> -> f32 + %adec1 = scalar.fmaf %av1, %cp1, %ag0_1 : f32 + %vsub1 = scalar.subf %vval1, %adec1 : f32 + %active_index1 = index.add %t0, %k1 : index + %row_in_range1 = index.cmp ult, %active_index1, %n_tokens : index + %active1 = scalar.ori %row_in_range1, %all_recurrence_rows : i1 + %dl1 = scf.if %active1 -> (f32) { + %value = scalar.mulf %vsub1, %bv1 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval1, %bv1 : f32 + scf.yield %padded : f32 + } + %dh1 = scalar.mulf %dl1, %ev1 : f32 + %grw1_0_a = vector.load %gr_flat[%k16] : view<256xf32> -> vector<4xf32> + %grw1_0_bi = index.constant 20 : index + %grw1_0_b = vector.load %gr_flat[%grw1_0_bi] : view<256xf32> -> vector<4xf32> + %grw1_1_a = vector.load %gr_flat[%k24] : view<256xf32> -> vector<4xf32> + %grw1_1_bi = index.constant 28 : index + %grw1_1_b = vector.load %gr_flat[%grw1_1_bi] : view<256xf32> -> vector<4xf32> + %qkw1_0_a = vector.load %qkt_flat[%k16] : view<256xf32> -> vector<4xf32> + %qkw1_0_bi = index.constant 20 : index + %qkw1_0_b = vector.load %qkt_flat[%qkw1_0_bi] : view<256xf32> -> vector<4xf32> + %qkw1_1_a = vector.load %qkt_flat[%k24] : view<256xf32> -> vector<4xf32> + %qkw1_1_bi = index.constant 28 : index + %qkw1_1_b = vector.load %qkt_flat[%qkw1_1_bi] : view<256xf32> -> vector<4xf32> + %qkd1_1 = vector.extract %qkw1_0_a[1] : vector<4xf32> -> f32 + %qsum1 = scalar.fmaf %qkd1_1, %dl1, %aq0_1 : f32 + %ao1_r = scalar.fmaf %bvv1, %cp1, %qsum1 : f32 + %ao1 = scf.if %active1 -> (f32) { + %value = scalar.mulf %ao1_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi1_a = index.add %rec_col, %k128 : index + %aoi1 = index.assume %aoi1_a [range(%aoi1_a, 0, 2047)] : index + %ddh1 = scalar.fptrunc %dh1 : f32 to f16 + %publish_active1 = scalar.andi %lane_live, %active1 : i1 + scf.if %publish_active1 { + view.store %ao1, %aost_flat[%aoi1] : f32, view<2048xf32> + } + %gp1_2 = vector.extract %grw1_0_a[2] : vector<4xf32> -> f32 + %ag1_2 = scalar.fmaf %gp1_2, %dl1, %ag0_2 : f32 + %qp1_2 = vector.extract %qkw1_0_a[2] : vector<4xf32> -> f32 + %aq1_2 = scalar.fmaf %qp1_2, %dl1, %aq0_2 : f32 + %gp1_3 = vector.extract %grw1_0_a[3] : vector<4xf32> -> f32 + %ag1_3 = scalar.fmaf %gp1_3, %dl1, %ag0_3 : f32 + %qp1_3 = vector.extract %qkw1_0_a[3] : vector<4xf32> -> f32 + %aq1_3 = scalar.fmaf %qp1_3, %dl1, %aq0_3 : f32 + %gp1_4 = vector.extract %grw1_0_b[0] : vector<4xf32> -> f32 + %ag1_4 = scalar.fmaf %gp1_4, %dl1, %ag0_4 : f32 + %qp1_4 = vector.extract %qkw1_0_b[0] : vector<4xf32> -> f32 + %aq1_4 = scalar.fmaf %qp1_4, %dl1, %aq0_4 : f32 + %gp1_5 = vector.extract %grw1_0_b[1] : vector<4xf32> -> f32 + %ag1_5 = scalar.fmaf %gp1_5, %dl1, %ag0_5 : f32 + %qp1_5 = vector.extract %qkw1_0_b[1] : vector<4xf32> -> f32 + %aq1_5 = scalar.fmaf %qp1_5, %dl1, %aq0_5 : f32 + %gp1_6 = vector.extract %grw1_0_b[2] : vector<4xf32> -> f32 + %ag1_6 = scalar.fmaf %gp1_6, %dl1, %ag0_6 : f32 + %qp1_6 = vector.extract %qkw1_0_b[2] : vector<4xf32> -> f32 + %aq1_6 = scalar.fmaf %qp1_6, %dl1, %aq0_6 : f32 + %gp1_7 = vector.extract %grw1_0_b[3] : vector<4xf32> -> f32 + %ag1_7 = scalar.fmaf %gp1_7, %dl1, %ag0_7 : f32 + %qp1_7 = vector.extract %qkw1_0_b[3] : vector<4xf32> -> f32 + %aq1_7 = scalar.fmaf %qp1_7, %dl1, %aq0_7 : f32 + %gp1_8 = vector.extract %grw1_1_a[0] : vector<4xf32> -> f32 + %ag1_8 = scalar.fmaf %gp1_8, %dl1, %ag0_8 : f32 + %qp1_8 = vector.extract %qkw1_1_a[0] : vector<4xf32> -> f32 + %aq1_8 = scalar.fmaf %qp1_8, %dl1, %aq0_8 : f32 + %gp1_9 = vector.extract %grw1_1_a[1] : vector<4xf32> -> f32 + %ag1_9 = scalar.fmaf %gp1_9, %dl1, %ag0_9 : f32 + %qp1_9 = vector.extract %qkw1_1_a[1] : vector<4xf32> -> f32 + %aq1_9 = scalar.fmaf %qp1_9, %dl1, %aq0_9 : f32 + %gp1_10 = vector.extract %grw1_1_a[2] : vector<4xf32> -> f32 + %ag1_10 = scalar.fmaf %gp1_10, %dl1, %ag0_10 : f32 + %qp1_10 = vector.extract %qkw1_1_a[2] : vector<4xf32> -> f32 + %aq1_10 = scalar.fmaf %qp1_10, %dl1, %aq0_10 : f32 + %gp1_11 = vector.extract %grw1_1_a[3] : vector<4xf32> -> f32 + %ag1_11 = scalar.fmaf %gp1_11, %dl1, %ag0_11 : f32 + %qp1_11 = vector.extract %qkw1_1_a[3] : vector<4xf32> -> f32 + %aq1_11 = scalar.fmaf %qp1_11, %dl1, %aq0_11 : f32 + %gp1_12 = vector.extract %grw1_1_b[0] : vector<4xf32> -> f32 + %ag1_12 = scalar.fmaf %gp1_12, %dl1, %ag0_12 : f32 + %qp1_12 = vector.extract %qkw1_1_b[0] : vector<4xf32> -> f32 + %aq1_12 = scalar.fmaf %qp1_12, %dl1, %aq0_12 : f32 + %gp1_13 = vector.extract %grw1_1_b[1] : vector<4xf32> -> f32 + %ag1_13 = scalar.fmaf %gp1_13, %dl1, %ag0_13 : f32 + %qp1_13 = vector.extract %qkw1_1_b[1] : vector<4xf32> -> f32 + %aq1_13 = scalar.fmaf %qp1_13, %dl1, %aq0_13 : f32 + %gp1_14 = vector.extract %grw1_1_b[2] : vector<4xf32> -> f32 + %ag1_14 = scalar.fmaf %gp1_14, %dl1, %ag0_14 : f32 + %qp1_14 = vector.extract %qkw1_1_b[2] : vector<4xf32> -> f32 + %aq1_14 = scalar.fmaf %qp1_14, %dl1, %aq0_14 : f32 + %gp1_15 = vector.extract %grw1_1_b[3] : vector<4xf32> -> f32 + %ag1_15 = scalar.fmaf %gp1_15, %dl1, %ag0_15 : f32 + %qp1_15 = vector.extract %qkw1_1_b[3] : vector<4xf32> -> f32 + %aq1_15 = scalar.fmaf %qp1_15, %dl1, %aq0_15 : f32 + %gq2 = vector.load %gs_flat[%k8] : view<80xf32> -> vector<4xf32> + %bv2 = vector.extract %gq2[0] : vector<4xf32> -> f32 + %cp2 = vector.extract %gq2[1] : vector<4xf32> -> f32 + %ev2 = vector.extract %gq2[2] : vector<4xf32> -> f32 + %tk2_r = index.add %t0, %k2 : index + %tk2_ok = index.cmp ult, %tk2_r, %n_tokens : index + %tk2_s = scf.select %tk2_ok, %tk2_r, %c0 : index + %tk2 = index.assume %tk2_s [range(%tk2_s, 0, 1048575)] : index + %av2_i = index.add %tile_row, %k2 : index + %av2_b = index.assume %av2_i [range(%av2_i, 0, 271)] : index + %av2 = view.load %a_flat[%av2_b] : view<272xf32> -> f32 + %bvv2 = view.load %b_flat[%av2_b] : view<272xf32> -> f32 + %vrd2_a = index.add %rec_col, %k256 : index + %vrd2 = index.assume %vrd2_a [range(%vrd2_a, 0, 2047)] : index + %vval2 = view.load %vst_flat[%vrd2] : view<2048xf32> -> f32 + %adec2 = scalar.fmaf %av2, %cp2, %ag1_2 : f32 + %vsub2 = scalar.subf %vval2, %adec2 : f32 + %active_index2 = index.add %t0, %k2 : index + %row_in_range2 = index.cmp ult, %active_index2, %n_tokens : index + %active2 = scalar.ori %row_in_range2, %all_recurrence_rows : i1 + %dl2 = scf.if %active2 -> (f32) { + %value = scalar.mulf %vsub2, %bv2 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval2, %bv2 : f32 + scf.yield %padded : f32 + } + %dh2 = scalar.mulf %dl2, %ev2 : f32 + %grw2_0_a = vector.load %gr_flat[%k32] : view<256xf32> -> vector<4xf32> + %grw2_0_bi = index.constant 36 : index + %grw2_0_b = vector.load %gr_flat[%grw2_0_bi] : view<256xf32> -> vector<4xf32> + %grw2_1_a = vector.load %gr_flat[%k40] : view<256xf32> -> vector<4xf32> + %grw2_1_bi = index.constant 44 : index + %grw2_1_b = vector.load %gr_flat[%grw2_1_bi] : view<256xf32> -> vector<4xf32> + %qkw2_0_a = vector.load %qkt_flat[%k32] : view<256xf32> -> vector<4xf32> + %qkw2_0_bi = index.constant 36 : index + %qkw2_0_b = vector.load %qkt_flat[%qkw2_0_bi] : view<256xf32> -> vector<4xf32> + %qkw2_1_a = vector.load %qkt_flat[%k40] : view<256xf32> -> vector<4xf32> + %qkw2_1_bi = index.constant 44 : index + %qkw2_1_b = vector.load %qkt_flat[%qkw2_1_bi] : view<256xf32> -> vector<4xf32> + %qkd2_2 = vector.extract %qkw2_0_a[2] : vector<4xf32> -> f32 + %qsum2 = scalar.fmaf %qkd2_2, %dl2, %aq1_2 : f32 + %ao2_r = scalar.fmaf %bvv2, %cp2, %qsum2 : f32 + %ao2 = scf.if %active2 -> (f32) { + %value = scalar.mulf %ao2_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi2_a = index.add %rec_col, %k256 : index + %aoi2 = index.assume %aoi2_a [range(%aoi2_a, 0, 2047)] : index + %ddh2 = scalar.fptrunc %dh2 : f32 to f16 + %publish_active2 = scalar.andi %lane_live, %active2 : i1 + scf.if %publish_active2 { + view.store %ao2, %aost_flat[%aoi2] : f32, view<2048xf32> + } + %gp2_3 = vector.extract %grw2_0_a[3] : vector<4xf32> -> f32 + %ag2_3 = scalar.fmaf %gp2_3, %dl2, %ag1_3 : f32 + %qp2_3 = vector.extract %qkw2_0_a[3] : vector<4xf32> -> f32 + %aq2_3 = scalar.fmaf %qp2_3, %dl2, %aq1_3 : f32 + %gp2_4 = vector.extract %grw2_0_b[0] : vector<4xf32> -> f32 + %ag2_4 = scalar.fmaf %gp2_4, %dl2, %ag1_4 : f32 + %qp2_4 = vector.extract %qkw2_0_b[0] : vector<4xf32> -> f32 + %aq2_4 = scalar.fmaf %qp2_4, %dl2, %aq1_4 : f32 + %gp2_5 = vector.extract %grw2_0_b[1] : vector<4xf32> -> f32 + %ag2_5 = scalar.fmaf %gp2_5, %dl2, %ag1_5 : f32 + %qp2_5 = vector.extract %qkw2_0_b[1] : vector<4xf32> -> f32 + %aq2_5 = scalar.fmaf %qp2_5, %dl2, %aq1_5 : f32 + %gp2_6 = vector.extract %grw2_0_b[2] : vector<4xf32> -> f32 + %ag2_6 = scalar.fmaf %gp2_6, %dl2, %ag1_6 : f32 + %qp2_6 = vector.extract %qkw2_0_b[2] : vector<4xf32> -> f32 + %aq2_6 = scalar.fmaf %qp2_6, %dl2, %aq1_6 : f32 + %gp2_7 = vector.extract %grw2_0_b[3] : vector<4xf32> -> f32 + %ag2_7 = scalar.fmaf %gp2_7, %dl2, %ag1_7 : f32 + %qp2_7 = vector.extract %qkw2_0_b[3] : vector<4xf32> -> f32 + %aq2_7 = scalar.fmaf %qp2_7, %dl2, %aq1_7 : f32 + %gp2_8 = vector.extract %grw2_1_a[0] : vector<4xf32> -> f32 + %ag2_8 = scalar.fmaf %gp2_8, %dl2, %ag1_8 : f32 + %qp2_8 = vector.extract %qkw2_1_a[0] : vector<4xf32> -> f32 + %aq2_8 = scalar.fmaf %qp2_8, %dl2, %aq1_8 : f32 + %gp2_9 = vector.extract %grw2_1_a[1] : vector<4xf32> -> f32 + %ag2_9 = scalar.fmaf %gp2_9, %dl2, %ag1_9 : f32 + %qp2_9 = vector.extract %qkw2_1_a[1] : vector<4xf32> -> f32 + %aq2_9 = scalar.fmaf %qp2_9, %dl2, %aq1_9 : f32 + %gp2_10 = vector.extract %grw2_1_a[2] : vector<4xf32> -> f32 + %ag2_10 = scalar.fmaf %gp2_10, %dl2, %ag1_10 : f32 + %qp2_10 = vector.extract %qkw2_1_a[2] : vector<4xf32> -> f32 + %aq2_10 = scalar.fmaf %qp2_10, %dl2, %aq1_10 : f32 + %gp2_11 = vector.extract %grw2_1_a[3] : vector<4xf32> -> f32 + %ag2_11 = scalar.fmaf %gp2_11, %dl2, %ag1_11 : f32 + %qp2_11 = vector.extract %qkw2_1_a[3] : vector<4xf32> -> f32 + %aq2_11 = scalar.fmaf %qp2_11, %dl2, %aq1_11 : f32 + %gp2_12 = vector.extract %grw2_1_b[0] : vector<4xf32> -> f32 + %ag2_12 = scalar.fmaf %gp2_12, %dl2, %ag1_12 : f32 + %qp2_12 = vector.extract %qkw2_1_b[0] : vector<4xf32> -> f32 + %aq2_12 = scalar.fmaf %qp2_12, %dl2, %aq1_12 : f32 + %gp2_13 = vector.extract %grw2_1_b[1] : vector<4xf32> -> f32 + %ag2_13 = scalar.fmaf %gp2_13, %dl2, %ag1_13 : f32 + %qp2_13 = vector.extract %qkw2_1_b[1] : vector<4xf32> -> f32 + %aq2_13 = scalar.fmaf %qp2_13, %dl2, %aq1_13 : f32 + %gp2_14 = vector.extract %grw2_1_b[2] : vector<4xf32> -> f32 + %ag2_14 = scalar.fmaf %gp2_14, %dl2, %ag1_14 : f32 + %qp2_14 = vector.extract %qkw2_1_b[2] : vector<4xf32> -> f32 + %aq2_14 = scalar.fmaf %qp2_14, %dl2, %aq1_14 : f32 + %gp2_15 = vector.extract %grw2_1_b[3] : vector<4xf32> -> f32 + %ag2_15 = scalar.fmaf %gp2_15, %dl2, %ag1_15 : f32 + %qp2_15 = vector.extract %qkw2_1_b[3] : vector<4xf32> -> f32 + %aq2_15 = scalar.fmaf %qp2_15, %dl2, %aq1_15 : f32 + %gq3 = vector.load %gs_flat[%k12] : view<80xf32> -> vector<4xf32> + %bv3 = vector.extract %gq3[0] : vector<4xf32> -> f32 + %cp3 = vector.extract %gq3[1] : vector<4xf32> -> f32 + %ev3 = vector.extract %gq3[2] : vector<4xf32> -> f32 + %tk3_r = index.add %t0, %k3 : index + %tk3_ok = index.cmp ult, %tk3_r, %n_tokens : index + %tk3_s = scf.select %tk3_ok, %tk3_r, %c0 : index + %tk3 = index.assume %tk3_s [range(%tk3_s, 0, 1048575)] : index + %av3_i = index.add %tile_row, %k3 : index + %av3_b = index.assume %av3_i [range(%av3_i, 0, 271)] : index + %av3 = view.load %a_flat[%av3_b] : view<272xf32> -> f32 + %bvv3 = view.load %b_flat[%av3_b] : view<272xf32> -> f32 + %vrd3_a = index.add %rec_col, %k384 : index + %vrd3 = index.assume %vrd3_a [range(%vrd3_a, 0, 2047)] : index + %vval3 = view.load %vst_flat[%vrd3] : view<2048xf32> -> f32 + %adec3 = scalar.fmaf %av3, %cp3, %ag2_3 : f32 + %vsub3 = scalar.subf %vval3, %adec3 : f32 + %active_index3 = index.add %t0, %k3 : index + %row_in_range3 = index.cmp ult, %active_index3, %n_tokens : index + %active3 = scalar.ori %row_in_range3, %all_recurrence_rows : i1 + %dl3 = scf.if %active3 -> (f32) { + %value = scalar.mulf %vsub3, %bv3 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval3, %bv3 : f32 + scf.yield %padded : f32 + } + %dh3 = scalar.mulf %dl3, %ev3 : f32 + %grw3_0_a = vector.load %gr_flat[%k48] : view<256xf32> -> vector<4xf32> + %grw3_0_bi = index.constant 52 : index + %grw3_0_b = vector.load %gr_flat[%grw3_0_bi] : view<256xf32> -> vector<4xf32> + %grw3_1_a = vector.load %gr_flat[%k56] : view<256xf32> -> vector<4xf32> + %grw3_1_bi = index.constant 60 : index + %grw3_1_b = vector.load %gr_flat[%grw3_1_bi] : view<256xf32> -> vector<4xf32> + %qkw3_0_a = vector.load %qkt_flat[%k48] : view<256xf32> -> vector<4xf32> + %qkw3_0_bi = index.constant 52 : index + %qkw3_0_b = vector.load %qkt_flat[%qkw3_0_bi] : view<256xf32> -> vector<4xf32> + %qkw3_1_a = vector.load %qkt_flat[%k56] : view<256xf32> -> vector<4xf32> + %qkw3_1_bi = index.constant 60 : index + %qkw3_1_b = vector.load %qkt_flat[%qkw3_1_bi] : view<256xf32> -> vector<4xf32> + %qkd3_3 = vector.extract %qkw3_0_a[3] : vector<4xf32> -> f32 + %qsum3 = scalar.fmaf %qkd3_3, %dl3, %aq2_3 : f32 + %ao3_r = scalar.fmaf %bvv3, %cp3, %qsum3 : f32 + %ao3 = scf.if %active3 -> (f32) { + %value = scalar.mulf %ao3_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi3_a = index.add %rec_col, %k384 : index + %aoi3 = index.assume %aoi3_a [range(%aoi3_a, 0, 2047)] : index + %ddh3 = scalar.fptrunc %dh3 : f32 to f16 + %publish_active3 = scalar.andi %lane_live, %active3 : i1 + scf.if %publish_active3 { + view.store %ao3, %aost_flat[%aoi3] : f32, view<2048xf32> + } + %gp3_4 = vector.extract %grw3_0_b[0] : vector<4xf32> -> f32 + %ag3_4 = scalar.fmaf %gp3_4, %dl3, %ag2_4 : f32 + %qp3_4 = vector.extract %qkw3_0_b[0] : vector<4xf32> -> f32 + %aq3_4 = scalar.fmaf %qp3_4, %dl3, %aq2_4 : f32 + %gp3_5 = vector.extract %grw3_0_b[1] : vector<4xf32> -> f32 + %ag3_5 = scalar.fmaf %gp3_5, %dl3, %ag2_5 : f32 + %qp3_5 = vector.extract %qkw3_0_b[1] : vector<4xf32> -> f32 + %aq3_5 = scalar.fmaf %qp3_5, %dl3, %aq2_5 : f32 + %gp3_6 = vector.extract %grw3_0_b[2] : vector<4xf32> -> f32 + %ag3_6 = scalar.fmaf %gp3_6, %dl3, %ag2_6 : f32 + %qp3_6 = vector.extract %qkw3_0_b[2] : vector<4xf32> -> f32 + %aq3_6 = scalar.fmaf %qp3_6, %dl3, %aq2_6 : f32 + %gp3_7 = vector.extract %grw3_0_b[3] : vector<4xf32> -> f32 + %ag3_7 = scalar.fmaf %gp3_7, %dl3, %ag2_7 : f32 + %qp3_7 = vector.extract %qkw3_0_b[3] : vector<4xf32> -> f32 + %aq3_7 = scalar.fmaf %qp3_7, %dl3, %aq2_7 : f32 + %gp3_8 = vector.extract %grw3_1_a[0] : vector<4xf32> -> f32 + %ag3_8 = scalar.fmaf %gp3_8, %dl3, %ag2_8 : f32 + %qp3_8 = vector.extract %qkw3_1_a[0] : vector<4xf32> -> f32 + %aq3_8 = scalar.fmaf %qp3_8, %dl3, %aq2_8 : f32 + %gp3_9 = vector.extract %grw3_1_a[1] : vector<4xf32> -> f32 + %ag3_9 = scalar.fmaf %gp3_9, %dl3, %ag2_9 : f32 + %qp3_9 = vector.extract %qkw3_1_a[1] : vector<4xf32> -> f32 + %aq3_9 = scalar.fmaf %qp3_9, %dl3, %aq2_9 : f32 + %gp3_10 = vector.extract %grw3_1_a[2] : vector<4xf32> -> f32 + %ag3_10 = scalar.fmaf %gp3_10, %dl3, %ag2_10 : f32 + %qp3_10 = vector.extract %qkw3_1_a[2] : vector<4xf32> -> f32 + %aq3_10 = scalar.fmaf %qp3_10, %dl3, %aq2_10 : f32 + %gp3_11 = vector.extract %grw3_1_a[3] : vector<4xf32> -> f32 + %ag3_11 = scalar.fmaf %gp3_11, %dl3, %ag2_11 : f32 + %qp3_11 = vector.extract %qkw3_1_a[3] : vector<4xf32> -> f32 + %aq3_11 = scalar.fmaf %qp3_11, %dl3, %aq2_11 : f32 + %gp3_12 = vector.extract %grw3_1_b[0] : vector<4xf32> -> f32 + %ag3_12 = scalar.fmaf %gp3_12, %dl3, %ag2_12 : f32 + %qp3_12 = vector.extract %qkw3_1_b[0] : vector<4xf32> -> f32 + %aq3_12 = scalar.fmaf %qp3_12, %dl3, %aq2_12 : f32 + %gp3_13 = vector.extract %grw3_1_b[1] : vector<4xf32> -> f32 + %ag3_13 = scalar.fmaf %gp3_13, %dl3, %ag2_13 : f32 + %qp3_13 = vector.extract %qkw3_1_b[1] : vector<4xf32> -> f32 + %aq3_13 = scalar.fmaf %qp3_13, %dl3, %aq2_13 : f32 + %gp3_14 = vector.extract %grw3_1_b[2] : vector<4xf32> -> f32 + %ag3_14 = scalar.fmaf %gp3_14, %dl3, %ag2_14 : f32 + %qp3_14 = vector.extract %qkw3_1_b[2] : vector<4xf32> -> f32 + %aq3_14 = scalar.fmaf %qp3_14, %dl3, %aq2_14 : f32 + %gp3_15 = vector.extract %grw3_1_b[3] : vector<4xf32> -> f32 + %ag3_15 = scalar.fmaf %gp3_15, %dl3, %ag2_15 : f32 + %qp3_15 = vector.extract %qkw3_1_b[3] : vector<4xf32> -> f32 + %aq3_15 = scalar.fmaf %qp3_15, %dl3, %aq2_15 : f32 + %gq4 = vector.load %gs_flat[%k16] : view<80xf32> -> vector<4xf32> + %bv4 = vector.extract %gq4[0] : vector<4xf32> -> f32 + %cp4 = vector.extract %gq4[1] : vector<4xf32> -> f32 + %ev4 = vector.extract %gq4[2] : vector<4xf32> -> f32 + %tk4_r = index.add %t0, %k4 : index + %tk4_ok = index.cmp ult, %tk4_r, %n_tokens : index + %tk4_s = scf.select %tk4_ok, %tk4_r, %c0 : index + %tk4 = index.assume %tk4_s [range(%tk4_s, 0, 1048575)] : index + %av4_i = index.add %tile_row, %k4 : index + %av4_b = index.assume %av4_i [range(%av4_i, 0, 271)] : index + %av4 = view.load %a_flat[%av4_b] : view<272xf32> -> f32 + %bvv4 = view.load %b_flat[%av4_b] : view<272xf32> -> f32 + %vrd4_a = index.add %rec_col, %k512 : index + %vrd4 = index.assume %vrd4_a [range(%vrd4_a, 0, 2047)] : index + %vval4 = view.load %vst_flat[%vrd4] : view<2048xf32> -> f32 + %adec4 = scalar.fmaf %av4, %cp4, %ag3_4 : f32 + %vsub4 = scalar.subf %vval4, %adec4 : f32 + %active_index4 = index.add %t0, %k4 : index + %row_in_range4 = index.cmp ult, %active_index4, %n_tokens : index + %active4 = scalar.ori %row_in_range4, %all_recurrence_rows : i1 + %dl4 = scf.if %active4 -> (f32) { + %value = scalar.mulf %vsub4, %bv4 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval4, %bv4 : f32 + scf.yield %padded : f32 + } + %dh4 = scalar.mulf %dl4, %ev4 : f32 + %grw4_0_a = vector.load %gr_flat[%k64] : view<256xf32> -> vector<4xf32> + %grw4_0_bi = index.constant 68 : index + %grw4_0_b = vector.load %gr_flat[%grw4_0_bi] : view<256xf32> -> vector<4xf32> + %grw4_1_a = vector.load %gr_flat[%k72] : view<256xf32> -> vector<4xf32> + %grw4_1_bi = index.constant 76 : index + %grw4_1_b = vector.load %gr_flat[%grw4_1_bi] : view<256xf32> -> vector<4xf32> + %qkw4_0_a = vector.load %qkt_flat[%k64] : view<256xf32> -> vector<4xf32> + %qkw4_0_bi = index.constant 68 : index + %qkw4_0_b = vector.load %qkt_flat[%qkw4_0_bi] : view<256xf32> -> vector<4xf32> + %qkw4_1_a = vector.load %qkt_flat[%k72] : view<256xf32> -> vector<4xf32> + %qkw4_1_bi = index.constant 76 : index + %qkw4_1_b = vector.load %qkt_flat[%qkw4_1_bi] : view<256xf32> -> vector<4xf32> + %qkd4_4 = vector.extract %qkw4_0_b[0] : vector<4xf32> -> f32 + %qsum4 = scalar.fmaf %qkd4_4, %dl4, %aq3_4 : f32 + %ao4_r = scalar.fmaf %bvv4, %cp4, %qsum4 : f32 + %ao4 = scf.if %active4 -> (f32) { + %value = scalar.mulf %ao4_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi4_a = index.add %rec_col, %k512 : index + %aoi4 = index.assume %aoi4_a [range(%aoi4_a, 0, 2047)] : index + %ddh4 = scalar.fptrunc %dh4 : f32 to f16 + %publish_active4 = scalar.andi %lane_live, %active4 : i1 + scf.if %publish_active4 { + view.store %ao4, %aost_flat[%aoi4] : f32, view<2048xf32> + } + %gp4_5 = vector.extract %grw4_0_b[1] : vector<4xf32> -> f32 + %ag4_5 = scalar.fmaf %gp4_5, %dl4, %ag3_5 : f32 + %qp4_5 = vector.extract %qkw4_0_b[1] : vector<4xf32> -> f32 + %aq4_5 = scalar.fmaf %qp4_5, %dl4, %aq3_5 : f32 + %gp4_6 = vector.extract %grw4_0_b[2] : vector<4xf32> -> f32 + %ag4_6 = scalar.fmaf %gp4_6, %dl4, %ag3_6 : f32 + %qp4_6 = vector.extract %qkw4_0_b[2] : vector<4xf32> -> f32 + %aq4_6 = scalar.fmaf %qp4_6, %dl4, %aq3_6 : f32 + %gp4_7 = vector.extract %grw4_0_b[3] : vector<4xf32> -> f32 + %ag4_7 = scalar.fmaf %gp4_7, %dl4, %ag3_7 : f32 + %qp4_7 = vector.extract %qkw4_0_b[3] : vector<4xf32> -> f32 + %aq4_7 = scalar.fmaf %qp4_7, %dl4, %aq3_7 : f32 + %gp4_8 = vector.extract %grw4_1_a[0] : vector<4xf32> -> f32 + %ag4_8 = scalar.fmaf %gp4_8, %dl4, %ag3_8 : f32 + %qp4_8 = vector.extract %qkw4_1_a[0] : vector<4xf32> -> f32 + %aq4_8 = scalar.fmaf %qp4_8, %dl4, %aq3_8 : f32 + %gp4_9 = vector.extract %grw4_1_a[1] : vector<4xf32> -> f32 + %ag4_9 = scalar.fmaf %gp4_9, %dl4, %ag3_9 : f32 + %qp4_9 = vector.extract %qkw4_1_a[1] : vector<4xf32> -> f32 + %aq4_9 = scalar.fmaf %qp4_9, %dl4, %aq3_9 : f32 + %gp4_10 = vector.extract %grw4_1_a[2] : vector<4xf32> -> f32 + %ag4_10 = scalar.fmaf %gp4_10, %dl4, %ag3_10 : f32 + %qp4_10 = vector.extract %qkw4_1_a[2] : vector<4xf32> -> f32 + %aq4_10 = scalar.fmaf %qp4_10, %dl4, %aq3_10 : f32 + %gp4_11 = vector.extract %grw4_1_a[3] : vector<4xf32> -> f32 + %ag4_11 = scalar.fmaf %gp4_11, %dl4, %ag3_11 : f32 + %qp4_11 = vector.extract %qkw4_1_a[3] : vector<4xf32> -> f32 + %aq4_11 = scalar.fmaf %qp4_11, %dl4, %aq3_11 : f32 + %gp4_12 = vector.extract %grw4_1_b[0] : vector<4xf32> -> f32 + %ag4_12 = scalar.fmaf %gp4_12, %dl4, %ag3_12 : f32 + %qp4_12 = vector.extract %qkw4_1_b[0] : vector<4xf32> -> f32 + %aq4_12 = scalar.fmaf %qp4_12, %dl4, %aq3_12 : f32 + %gp4_13 = vector.extract %grw4_1_b[1] : vector<4xf32> -> f32 + %ag4_13 = scalar.fmaf %gp4_13, %dl4, %ag3_13 : f32 + %qp4_13 = vector.extract %qkw4_1_b[1] : vector<4xf32> -> f32 + %aq4_13 = scalar.fmaf %qp4_13, %dl4, %aq3_13 : f32 + %gp4_14 = vector.extract %grw4_1_b[2] : vector<4xf32> -> f32 + %ag4_14 = scalar.fmaf %gp4_14, %dl4, %ag3_14 : f32 + %qp4_14 = vector.extract %qkw4_1_b[2] : vector<4xf32> -> f32 + %aq4_14 = scalar.fmaf %qp4_14, %dl4, %aq3_14 : f32 + %gp4_15 = vector.extract %grw4_1_b[3] : vector<4xf32> -> f32 + %ag4_15 = scalar.fmaf %gp4_15, %dl4, %ag3_15 : f32 + %qp4_15 = vector.extract %qkw4_1_b[3] : vector<4xf32> -> f32 + %aq4_15 = scalar.fmaf %qp4_15, %dl4, %aq3_15 : f32 + %gq5 = vector.load %gs_flat[%k20] : view<80xf32> -> vector<4xf32> + %bv5 = vector.extract %gq5[0] : vector<4xf32> -> f32 + %cp5 = vector.extract %gq5[1] : vector<4xf32> -> f32 + %ev5 = vector.extract %gq5[2] : vector<4xf32> -> f32 + %tk5_r = index.add %t0, %k5 : index + %tk5_ok = index.cmp ult, %tk5_r, %n_tokens : index + %tk5_s = scf.select %tk5_ok, %tk5_r, %c0 : index + %tk5 = index.assume %tk5_s [range(%tk5_s, 0, 1048575)] : index + %av5_i = index.add %tile_row, %k5 : index + %av5_b = index.assume %av5_i [range(%av5_i, 0, 271)] : index + %av5 = view.load %a_flat[%av5_b] : view<272xf32> -> f32 + %bvv5 = view.load %b_flat[%av5_b] : view<272xf32> -> f32 + %vrd5_a = index.add %rec_col, %k640 : index + %vrd5 = index.assume %vrd5_a [range(%vrd5_a, 0, 2047)] : index + %vval5 = view.load %vst_flat[%vrd5] : view<2048xf32> -> f32 + %adec5 = scalar.fmaf %av5, %cp5, %ag4_5 : f32 + %vsub5 = scalar.subf %vval5, %adec5 : f32 + %active_index5 = index.add %t0, %k5 : index + %row_in_range5 = index.cmp ult, %active_index5, %n_tokens : index + %active5 = scalar.ori %row_in_range5, %all_recurrence_rows : i1 + %dl5 = scf.if %active5 -> (f32) { + %value = scalar.mulf %vsub5, %bv5 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval5, %bv5 : f32 + scf.yield %padded : f32 + } + %dh5 = scalar.mulf %dl5, %ev5 : f32 + %grw5_0_a = vector.load %gr_flat[%k80] : view<256xf32> -> vector<4xf32> + %grw5_0_bi = index.constant 84 : index + %grw5_0_b = vector.load %gr_flat[%grw5_0_bi] : view<256xf32> -> vector<4xf32> + %grw5_1_a = vector.load %gr_flat[%k88] : view<256xf32> -> vector<4xf32> + %grw5_1_bi = index.constant 92 : index + %grw5_1_b = vector.load %gr_flat[%grw5_1_bi] : view<256xf32> -> vector<4xf32> + %qkw5_0_a = vector.load %qkt_flat[%k80] : view<256xf32> -> vector<4xf32> + %qkw5_0_bi = index.constant 84 : index + %qkw5_0_b = vector.load %qkt_flat[%qkw5_0_bi] : view<256xf32> -> vector<4xf32> + %qkw5_1_a = vector.load %qkt_flat[%k88] : view<256xf32> -> vector<4xf32> + %qkw5_1_bi = index.constant 92 : index + %qkw5_1_b = vector.load %qkt_flat[%qkw5_1_bi] : view<256xf32> -> vector<4xf32> + %qkd5_5 = vector.extract %qkw5_0_b[1] : vector<4xf32> -> f32 + %qsum5 = scalar.fmaf %qkd5_5, %dl5, %aq4_5 : f32 + %ao5_r = scalar.fmaf %bvv5, %cp5, %qsum5 : f32 + %ao5 = scf.if %active5 -> (f32) { + %value = scalar.mulf %ao5_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi5_a = index.add %rec_col, %k640 : index + %aoi5 = index.assume %aoi5_a [range(%aoi5_a, 0, 2047)] : index + %ddh5 = scalar.fptrunc %dh5 : f32 to f16 + %publish_active5 = scalar.andi %lane_live, %active5 : i1 + scf.if %publish_active5 { + view.store %ao5, %aost_flat[%aoi5] : f32, view<2048xf32> + } + %gp5_6 = vector.extract %grw5_0_b[2] : vector<4xf32> -> f32 + %ag5_6 = scalar.fmaf %gp5_6, %dl5, %ag4_6 : f32 + %qp5_6 = vector.extract %qkw5_0_b[2] : vector<4xf32> -> f32 + %aq5_6 = scalar.fmaf %qp5_6, %dl5, %aq4_6 : f32 + %gp5_7 = vector.extract %grw5_0_b[3] : vector<4xf32> -> f32 + %ag5_7 = scalar.fmaf %gp5_7, %dl5, %ag4_7 : f32 + %qp5_7 = vector.extract %qkw5_0_b[3] : vector<4xf32> -> f32 + %aq5_7 = scalar.fmaf %qp5_7, %dl5, %aq4_7 : f32 + %gp5_8 = vector.extract %grw5_1_a[0] : vector<4xf32> -> f32 + %ag5_8 = scalar.fmaf %gp5_8, %dl5, %ag4_8 : f32 + %qp5_8 = vector.extract %qkw5_1_a[0] : vector<4xf32> -> f32 + %aq5_8 = scalar.fmaf %qp5_8, %dl5, %aq4_8 : f32 + %gp5_9 = vector.extract %grw5_1_a[1] : vector<4xf32> -> f32 + %ag5_9 = scalar.fmaf %gp5_9, %dl5, %ag4_9 : f32 + %qp5_9 = vector.extract %qkw5_1_a[1] : vector<4xf32> -> f32 + %aq5_9 = scalar.fmaf %qp5_9, %dl5, %aq4_9 : f32 + %gp5_10 = vector.extract %grw5_1_a[2] : vector<4xf32> -> f32 + %ag5_10 = scalar.fmaf %gp5_10, %dl5, %ag4_10 : f32 + %qp5_10 = vector.extract %qkw5_1_a[2] : vector<4xf32> -> f32 + %aq5_10 = scalar.fmaf %qp5_10, %dl5, %aq4_10 : f32 + %gp5_11 = vector.extract %grw5_1_a[3] : vector<4xf32> -> f32 + %ag5_11 = scalar.fmaf %gp5_11, %dl5, %ag4_11 : f32 + %qp5_11 = vector.extract %qkw5_1_a[3] : vector<4xf32> -> f32 + %aq5_11 = scalar.fmaf %qp5_11, %dl5, %aq4_11 : f32 + %gp5_12 = vector.extract %grw5_1_b[0] : vector<4xf32> -> f32 + %ag5_12 = scalar.fmaf %gp5_12, %dl5, %ag4_12 : f32 + %qp5_12 = vector.extract %qkw5_1_b[0] : vector<4xf32> -> f32 + %aq5_12 = scalar.fmaf %qp5_12, %dl5, %aq4_12 : f32 + %gp5_13 = vector.extract %grw5_1_b[1] : vector<4xf32> -> f32 + %ag5_13 = scalar.fmaf %gp5_13, %dl5, %ag4_13 : f32 + %qp5_13 = vector.extract %qkw5_1_b[1] : vector<4xf32> -> f32 + %aq5_13 = scalar.fmaf %qp5_13, %dl5, %aq4_13 : f32 + %gp5_14 = vector.extract %grw5_1_b[2] : vector<4xf32> -> f32 + %ag5_14 = scalar.fmaf %gp5_14, %dl5, %ag4_14 : f32 + %qp5_14 = vector.extract %qkw5_1_b[2] : vector<4xf32> -> f32 + %aq5_14 = scalar.fmaf %qp5_14, %dl5, %aq4_14 : f32 + %gp5_15 = vector.extract %grw5_1_b[3] : vector<4xf32> -> f32 + %ag5_15 = scalar.fmaf %gp5_15, %dl5, %ag4_15 : f32 + %qp5_15 = vector.extract %qkw5_1_b[3] : vector<4xf32> -> f32 + %aq5_15 = scalar.fmaf %qp5_15, %dl5, %aq4_15 : f32 + %gq6 = vector.load %gs_flat[%k24] : view<80xf32> -> vector<4xf32> + %bv6 = vector.extract %gq6[0] : vector<4xf32> -> f32 + %cp6 = vector.extract %gq6[1] : vector<4xf32> -> f32 + %ev6 = vector.extract %gq6[2] : vector<4xf32> -> f32 + %tk6_r = index.add %t0, %k6 : index + %tk6_ok = index.cmp ult, %tk6_r, %n_tokens : index + %tk6_s = scf.select %tk6_ok, %tk6_r, %c0 : index + %tk6 = index.assume %tk6_s [range(%tk6_s, 0, 1048575)] : index + %av6_i = index.add %tile_row, %k6 : index + %av6_b = index.assume %av6_i [range(%av6_i, 0, 271)] : index + %av6 = view.load %a_flat[%av6_b] : view<272xf32> -> f32 + %bvv6 = view.load %b_flat[%av6_b] : view<272xf32> -> f32 + %vrd6_a = index.add %rec_col, %k768 : index + %vrd6 = index.assume %vrd6_a [range(%vrd6_a, 0, 2047)] : index + %vval6 = view.load %vst_flat[%vrd6] : view<2048xf32> -> f32 + %adec6 = scalar.fmaf %av6, %cp6, %ag5_6 : f32 + %vsub6 = scalar.subf %vval6, %adec6 : f32 + %active_index6 = index.add %t0, %k6 : index + %row_in_range6 = index.cmp ult, %active_index6, %n_tokens : index + %active6 = scalar.ori %row_in_range6, %all_recurrence_rows : i1 + %dl6 = scf.if %active6 -> (f32) { + %value = scalar.mulf %vsub6, %bv6 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval6, %bv6 : f32 + scf.yield %padded : f32 + } + %dh6 = scalar.mulf %dl6, %ev6 : f32 + %grw6_0_a = vector.load %gr_flat[%k96] : view<256xf32> -> vector<4xf32> + %grw6_0_bi = index.constant 100 : index + %grw6_0_b = vector.load %gr_flat[%grw6_0_bi] : view<256xf32> -> vector<4xf32> + %grw6_1_a = vector.load %gr_flat[%k104] : view<256xf32> -> vector<4xf32> + %grw6_1_bi = index.constant 108 : index + %grw6_1_b = vector.load %gr_flat[%grw6_1_bi] : view<256xf32> -> vector<4xf32> + %qkw6_0_a = vector.load %qkt_flat[%k96] : view<256xf32> -> vector<4xf32> + %qkw6_0_bi = index.constant 100 : index + %qkw6_0_b = vector.load %qkt_flat[%qkw6_0_bi] : view<256xf32> -> vector<4xf32> + %qkw6_1_a = vector.load %qkt_flat[%k104] : view<256xf32> -> vector<4xf32> + %qkw6_1_bi = index.constant 108 : index + %qkw6_1_b = vector.load %qkt_flat[%qkw6_1_bi] : view<256xf32> -> vector<4xf32> + %qkd6_6 = vector.extract %qkw6_0_b[2] : vector<4xf32> -> f32 + %qsum6 = scalar.fmaf %qkd6_6, %dl6, %aq5_6 : f32 + %ao6_r = scalar.fmaf %bvv6, %cp6, %qsum6 : f32 + %ao6 = scf.if %active6 -> (f32) { + %value = scalar.mulf %ao6_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi6_a = index.add %rec_col, %k768 : index + %aoi6 = index.assume %aoi6_a [range(%aoi6_a, 0, 2047)] : index + %ddh6 = scalar.fptrunc %dh6 : f32 to f16 + %publish_active6 = scalar.andi %lane_live, %active6 : i1 + scf.if %publish_active6 { + view.store %ao6, %aost_flat[%aoi6] : f32, view<2048xf32> + } + %gp6_7 = vector.extract %grw6_0_b[3] : vector<4xf32> -> f32 + %ag6_7 = scalar.fmaf %gp6_7, %dl6, %ag5_7 : f32 + %qp6_7 = vector.extract %qkw6_0_b[3] : vector<4xf32> -> f32 + %aq6_7 = scalar.fmaf %qp6_7, %dl6, %aq5_7 : f32 + %gp6_8 = vector.extract %grw6_1_a[0] : vector<4xf32> -> f32 + %ag6_8 = scalar.fmaf %gp6_8, %dl6, %ag5_8 : f32 + %qp6_8 = vector.extract %qkw6_1_a[0] : vector<4xf32> -> f32 + %aq6_8 = scalar.fmaf %qp6_8, %dl6, %aq5_8 : f32 + %gp6_9 = vector.extract %grw6_1_a[1] : vector<4xf32> -> f32 + %ag6_9 = scalar.fmaf %gp6_9, %dl6, %ag5_9 : f32 + %qp6_9 = vector.extract %qkw6_1_a[1] : vector<4xf32> -> f32 + %aq6_9 = scalar.fmaf %qp6_9, %dl6, %aq5_9 : f32 + %gp6_10 = vector.extract %grw6_1_a[2] : vector<4xf32> -> f32 + %ag6_10 = scalar.fmaf %gp6_10, %dl6, %ag5_10 : f32 + %qp6_10 = vector.extract %qkw6_1_a[2] : vector<4xf32> -> f32 + %aq6_10 = scalar.fmaf %qp6_10, %dl6, %aq5_10 : f32 + %gp6_11 = vector.extract %grw6_1_a[3] : vector<4xf32> -> f32 + %ag6_11 = scalar.fmaf %gp6_11, %dl6, %ag5_11 : f32 + %qp6_11 = vector.extract %qkw6_1_a[3] : vector<4xf32> -> f32 + %aq6_11 = scalar.fmaf %qp6_11, %dl6, %aq5_11 : f32 + %gp6_12 = vector.extract %grw6_1_b[0] : vector<4xf32> -> f32 + %ag6_12 = scalar.fmaf %gp6_12, %dl6, %ag5_12 : f32 + %qp6_12 = vector.extract %qkw6_1_b[0] : vector<4xf32> -> f32 + %aq6_12 = scalar.fmaf %qp6_12, %dl6, %aq5_12 : f32 + %gp6_13 = vector.extract %grw6_1_b[1] : vector<4xf32> -> f32 + %ag6_13 = scalar.fmaf %gp6_13, %dl6, %ag5_13 : f32 + %qp6_13 = vector.extract %qkw6_1_b[1] : vector<4xf32> -> f32 + %aq6_13 = scalar.fmaf %qp6_13, %dl6, %aq5_13 : f32 + %gp6_14 = vector.extract %grw6_1_b[2] : vector<4xf32> -> f32 + %ag6_14 = scalar.fmaf %gp6_14, %dl6, %ag5_14 : f32 + %qp6_14 = vector.extract %qkw6_1_b[2] : vector<4xf32> -> f32 + %aq6_14 = scalar.fmaf %qp6_14, %dl6, %aq5_14 : f32 + %gp6_15 = vector.extract %grw6_1_b[3] : vector<4xf32> -> f32 + %ag6_15 = scalar.fmaf %gp6_15, %dl6, %ag5_15 : f32 + %qp6_15 = vector.extract %qkw6_1_b[3] : vector<4xf32> -> f32 + %aq6_15 = scalar.fmaf %qp6_15, %dl6, %aq5_15 : f32 + %gq7 = vector.load %gs_flat[%k28] : view<80xf32> -> vector<4xf32> + %bv7 = vector.extract %gq7[0] : vector<4xf32> -> f32 + %cp7 = vector.extract %gq7[1] : vector<4xf32> -> f32 + %ev7 = vector.extract %gq7[2] : vector<4xf32> -> f32 + %tk7_r = index.add %t0, %k7 : index + %tk7_ok = index.cmp ult, %tk7_r, %n_tokens : index + %tk7_s = scf.select %tk7_ok, %tk7_r, %c0 : index + %tk7 = index.assume %tk7_s [range(%tk7_s, 0, 1048575)] : index + %av7_i = index.add %tile_row, %k7 : index + %av7_b = index.assume %av7_i [range(%av7_i, 0, 271)] : index + %av7 = view.load %a_flat[%av7_b] : view<272xf32> -> f32 + %bvv7 = view.load %b_flat[%av7_b] : view<272xf32> -> f32 + %vrd7_a = index.add %rec_col, %k896 : index + %vrd7 = index.assume %vrd7_a [range(%vrd7_a, 0, 2047)] : index + %vval7 = view.load %vst_flat[%vrd7] : view<2048xf32> -> f32 + %adec7 = scalar.fmaf %av7, %cp7, %ag6_7 : f32 + %vsub7 = scalar.subf %vval7, %adec7 : f32 + %active_index7 = index.add %t0, %k7 : index + %row_in_range7 = index.cmp ult, %active_index7, %n_tokens : index + %active7 = scalar.ori %row_in_range7, %all_recurrence_rows : i1 + %dl7 = scf.if %active7 -> (f32) { + %value = scalar.mulf %vsub7, %bv7 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval7, %bv7 : f32 + scf.yield %padded : f32 + } + %dh7 = scalar.mulf %dl7, %ev7 : f32 + %grw7_0_a = vector.load %gr_flat[%k112] : view<256xf32> -> vector<4xf32> + %grw7_0_bi = index.constant 116 : index + %grw7_0_b = vector.load %gr_flat[%grw7_0_bi] : view<256xf32> -> vector<4xf32> + %grw7_1_a = vector.load %gr_flat[%k120] : view<256xf32> -> vector<4xf32> + %grw7_1_bi = index.constant 124 : index + %grw7_1_b = vector.load %gr_flat[%grw7_1_bi] : view<256xf32> -> vector<4xf32> + %qkw7_0_a = vector.load %qkt_flat[%k112] : view<256xf32> -> vector<4xf32> + %qkw7_0_bi = index.constant 116 : index + %qkw7_0_b = vector.load %qkt_flat[%qkw7_0_bi] : view<256xf32> -> vector<4xf32> + %qkw7_1_a = vector.load %qkt_flat[%k120] : view<256xf32> -> vector<4xf32> + %qkw7_1_bi = index.constant 124 : index + %qkw7_1_b = vector.load %qkt_flat[%qkw7_1_bi] : view<256xf32> -> vector<4xf32> + %qkd7_7 = vector.extract %qkw7_0_b[3] : vector<4xf32> -> f32 + %qsum7 = scalar.fmaf %qkd7_7, %dl7, %aq6_7 : f32 + %ao7_r = scalar.fmaf %bvv7, %cp7, %qsum7 : f32 + %ao7 = scf.if %active7 -> (f32) { + %value = scalar.mulf %ao7_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi7_a = index.add %rec_col, %k896 : index + %aoi7 = index.assume %aoi7_a [range(%aoi7_a, 0, 2047)] : index + %ddh7 = scalar.fptrunc %dh7 : f32 to f16 + %publish_active7 = scalar.andi %lane_live, %active7 : i1 + scf.if %publish_active7 { + view.store %ao7, %aost_flat[%aoi7] : f32, view<2048xf32> + } + %gp7_8 = vector.extract %grw7_1_a[0] : vector<4xf32> -> f32 + %ag7_8 = scalar.fmaf %gp7_8, %dl7, %ag6_8 : f32 + %qp7_8 = vector.extract %qkw7_1_a[0] : vector<4xf32> -> f32 + %aq7_8 = scalar.fmaf %qp7_8, %dl7, %aq6_8 : f32 + %gp7_9 = vector.extract %grw7_1_a[1] : vector<4xf32> -> f32 + %ag7_9 = scalar.fmaf %gp7_9, %dl7, %ag6_9 : f32 + %qp7_9 = vector.extract %qkw7_1_a[1] : vector<4xf32> -> f32 + %aq7_9 = scalar.fmaf %qp7_9, %dl7, %aq6_9 : f32 + %gp7_10 = vector.extract %grw7_1_a[2] : vector<4xf32> -> f32 + %ag7_10 = scalar.fmaf %gp7_10, %dl7, %ag6_10 : f32 + %qp7_10 = vector.extract %qkw7_1_a[2] : vector<4xf32> -> f32 + %aq7_10 = scalar.fmaf %qp7_10, %dl7, %aq6_10 : f32 + %gp7_11 = vector.extract %grw7_1_a[3] : vector<4xf32> -> f32 + %ag7_11 = scalar.fmaf %gp7_11, %dl7, %ag6_11 : f32 + %qp7_11 = vector.extract %qkw7_1_a[3] : vector<4xf32> -> f32 + %aq7_11 = scalar.fmaf %qp7_11, %dl7, %aq6_11 : f32 + %gp7_12 = vector.extract %grw7_1_b[0] : vector<4xf32> -> f32 + %ag7_12 = scalar.fmaf %gp7_12, %dl7, %ag6_12 : f32 + %qp7_12 = vector.extract %qkw7_1_b[0] : vector<4xf32> -> f32 + %aq7_12 = scalar.fmaf %qp7_12, %dl7, %aq6_12 : f32 + %gp7_13 = vector.extract %grw7_1_b[1] : vector<4xf32> -> f32 + %ag7_13 = scalar.fmaf %gp7_13, %dl7, %ag6_13 : f32 + %qp7_13 = vector.extract %qkw7_1_b[1] : vector<4xf32> -> f32 + %aq7_13 = scalar.fmaf %qp7_13, %dl7, %aq6_13 : f32 + %gp7_14 = vector.extract %grw7_1_b[2] : vector<4xf32> -> f32 + %ag7_14 = scalar.fmaf %gp7_14, %dl7, %ag6_14 : f32 + %qp7_14 = vector.extract %qkw7_1_b[2] : vector<4xf32> -> f32 + %aq7_14 = scalar.fmaf %qp7_14, %dl7, %aq6_14 : f32 + %gp7_15 = vector.extract %grw7_1_b[3] : vector<4xf32> -> f32 + %ag7_15 = scalar.fmaf %gp7_15, %dl7, %ag6_15 : f32 + %qp7_15 = vector.extract %qkw7_1_b[3] : vector<4xf32> -> f32 + %aq7_15 = scalar.fmaf %qp7_15, %dl7, %aq6_15 : f32 + %gq8 = vector.load %gs_flat[%k32] : view<80xf32> -> vector<4xf32> + %bv8 = vector.extract %gq8[0] : vector<4xf32> -> f32 + %cp8 = vector.extract %gq8[1] : vector<4xf32> -> f32 + %ev8 = vector.extract %gq8[2] : vector<4xf32> -> f32 + %tk8_r = index.add %t0, %k8 : index + %tk8_ok = index.cmp ult, %tk8_r, %n_tokens : index + %tk8_s = scf.select %tk8_ok, %tk8_r, %c0 : index + %tk8 = index.assume %tk8_s [range(%tk8_s, 0, 1048575)] : index + %av8_i = index.add %tile_row, %k8 : index + %av8_b = index.assume %av8_i [range(%av8_i, 0, 271)] : index + %av8 = view.load %a_flat[%av8_b] : view<272xf32> -> f32 + %bvv8 = view.load %b_flat[%av8_b] : view<272xf32> -> f32 + %vrd8_a = index.add %rec_col, %k1024 : index + %vrd8 = index.assume %vrd8_a [range(%vrd8_a, 0, 2047)] : index + %vval8 = view.load %vst_flat[%vrd8] : view<2048xf32> -> f32 + %adec8 = scalar.fmaf %av8, %cp8, %ag7_8 : f32 + %vsub8 = scalar.subf %vval8, %adec8 : f32 + %active_index8 = index.add %t0, %k8 : index + %row_in_range8 = index.cmp ult, %active_index8, %n_tokens : index + %active8 = scalar.ori %row_in_range8, %all_recurrence_rows : i1 + %dl8 = scf.if %active8 -> (f32) { + %value = scalar.mulf %vsub8, %bv8 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval8, %bv8 : f32 + scf.yield %padded : f32 + } + %dh8 = scalar.mulf %dl8, %ev8 : f32 + %grw8_0_a = vector.load %gr_flat[%k128] : view<256xf32> -> vector<4xf32> + %grw8_0_bi = index.constant 132 : index + %grw8_0_b = vector.load %gr_flat[%grw8_0_bi] : view<256xf32> -> vector<4xf32> + %grw8_1_a = vector.load %gr_flat[%k136] : view<256xf32> -> vector<4xf32> + %grw8_1_bi = index.constant 140 : index + %grw8_1_b = vector.load %gr_flat[%grw8_1_bi] : view<256xf32> -> vector<4xf32> + %qkw8_0_a = vector.load %qkt_flat[%k128] : view<256xf32> -> vector<4xf32> + %qkw8_0_bi = index.constant 132 : index + %qkw8_0_b = vector.load %qkt_flat[%qkw8_0_bi] : view<256xf32> -> vector<4xf32> + %qkw8_1_a = vector.load %qkt_flat[%k136] : view<256xf32> -> vector<4xf32> + %qkw8_1_bi = index.constant 140 : index + %qkw8_1_b = vector.load %qkt_flat[%qkw8_1_bi] : view<256xf32> -> vector<4xf32> + %qkd8_8 = vector.extract %qkw8_1_a[0] : vector<4xf32> -> f32 + %qsum8 = scalar.fmaf %qkd8_8, %dl8, %aq7_8 : f32 + %ao8_r = scalar.fmaf %bvv8, %cp8, %qsum8 : f32 + %ao8 = scf.if %active8 -> (f32) { + %value = scalar.mulf %ao8_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi8_a = index.add %rec_col, %k1024 : index + %aoi8 = index.assume %aoi8_a [range(%aoi8_a, 0, 2047)] : index + %ddh8 = scalar.fptrunc %dh8 : f32 to f16 + %publish_active8 = scalar.andi %lane_live, %active8 : i1 + scf.if %publish_active8 { + view.store %ao8, %aost_flat[%aoi8] : f32, view<2048xf32> + } + %gp8_9 = vector.extract %grw8_1_a[1] : vector<4xf32> -> f32 + %ag8_9 = scalar.fmaf %gp8_9, %dl8, %ag7_9 : f32 + %qp8_9 = vector.extract %qkw8_1_a[1] : vector<4xf32> -> f32 + %aq8_9 = scalar.fmaf %qp8_9, %dl8, %aq7_9 : f32 + %gp8_10 = vector.extract %grw8_1_a[2] : vector<4xf32> -> f32 + %ag8_10 = scalar.fmaf %gp8_10, %dl8, %ag7_10 : f32 + %qp8_10 = vector.extract %qkw8_1_a[2] : vector<4xf32> -> f32 + %aq8_10 = scalar.fmaf %qp8_10, %dl8, %aq7_10 : f32 + %gp8_11 = vector.extract %grw8_1_a[3] : vector<4xf32> -> f32 + %ag8_11 = scalar.fmaf %gp8_11, %dl8, %ag7_11 : f32 + %qp8_11 = vector.extract %qkw8_1_a[3] : vector<4xf32> -> f32 + %aq8_11 = scalar.fmaf %qp8_11, %dl8, %aq7_11 : f32 + %gp8_12 = vector.extract %grw8_1_b[0] : vector<4xf32> -> f32 + %ag8_12 = scalar.fmaf %gp8_12, %dl8, %ag7_12 : f32 + %qp8_12 = vector.extract %qkw8_1_b[0] : vector<4xf32> -> f32 + %aq8_12 = scalar.fmaf %qp8_12, %dl8, %aq7_12 : f32 + %gp8_13 = vector.extract %grw8_1_b[1] : vector<4xf32> -> f32 + %ag8_13 = scalar.fmaf %gp8_13, %dl8, %ag7_13 : f32 + %qp8_13 = vector.extract %qkw8_1_b[1] : vector<4xf32> -> f32 + %aq8_13 = scalar.fmaf %qp8_13, %dl8, %aq7_13 : f32 + %gp8_14 = vector.extract %grw8_1_b[2] : vector<4xf32> -> f32 + %ag8_14 = scalar.fmaf %gp8_14, %dl8, %ag7_14 : f32 + %qp8_14 = vector.extract %qkw8_1_b[2] : vector<4xf32> -> f32 + %aq8_14 = scalar.fmaf %qp8_14, %dl8, %aq7_14 : f32 + %gp8_15 = vector.extract %grw8_1_b[3] : vector<4xf32> -> f32 + %ag8_15 = scalar.fmaf %gp8_15, %dl8, %ag7_15 : f32 + %qp8_15 = vector.extract %qkw8_1_b[3] : vector<4xf32> -> f32 + %aq8_15 = scalar.fmaf %qp8_15, %dl8, %aq7_15 : f32 + %gq9 = vector.load %gs_flat[%k36] : view<80xf32> -> vector<4xf32> + %bv9 = vector.extract %gq9[0] : vector<4xf32> -> f32 + %cp9 = vector.extract %gq9[1] : vector<4xf32> -> f32 + %ev9 = vector.extract %gq9[2] : vector<4xf32> -> f32 + %tk9_r = index.add %t0, %k9 : index + %tk9_ok = index.cmp ult, %tk9_r, %n_tokens : index + %tk9_s = scf.select %tk9_ok, %tk9_r, %c0 : index + %tk9 = index.assume %tk9_s [range(%tk9_s, 0, 1048575)] : index + %av9_i = index.add %tile_row, %k9 : index + %av9_b = index.assume %av9_i [range(%av9_i, 0, 271)] : index + %av9 = view.load %a_flat[%av9_b] : view<272xf32> -> f32 + %bvv9 = view.load %b_flat[%av9_b] : view<272xf32> -> f32 + %vrd9_a = index.add %rec_col, %k1152 : index + %vrd9 = index.assume %vrd9_a [range(%vrd9_a, 0, 2047)] : index + %vval9 = view.load %vst_flat[%vrd9] : view<2048xf32> -> f32 + %adec9 = scalar.fmaf %av9, %cp9, %ag8_9 : f32 + %vsub9 = scalar.subf %vval9, %adec9 : f32 + %active_index9 = index.add %t0, %k9 : index + %row_in_range9 = index.cmp ult, %active_index9, %n_tokens : index + %active9 = scalar.ori %row_in_range9, %all_recurrence_rows : i1 + %dl9 = scf.if %active9 -> (f32) { + %value = scalar.mulf %vsub9, %bv9 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval9, %bv9 : f32 + scf.yield %padded : f32 + } + %dh9 = scalar.mulf %dl9, %ev9 : f32 + %grw9_0_a = vector.load %gr_flat[%k144] : view<256xf32> -> vector<4xf32> + %grw9_0_bi = index.constant 148 : index + %grw9_0_b = vector.load %gr_flat[%grw9_0_bi] : view<256xf32> -> vector<4xf32> + %grw9_1_a = vector.load %gr_flat[%k152] : view<256xf32> -> vector<4xf32> + %grw9_1_bi = index.constant 156 : index + %grw9_1_b = vector.load %gr_flat[%grw9_1_bi] : view<256xf32> -> vector<4xf32> + %qkw9_0_a = vector.load %qkt_flat[%k144] : view<256xf32> -> vector<4xf32> + %qkw9_0_bi = index.constant 148 : index + %qkw9_0_b = vector.load %qkt_flat[%qkw9_0_bi] : view<256xf32> -> vector<4xf32> + %qkw9_1_a = vector.load %qkt_flat[%k152] : view<256xf32> -> vector<4xf32> + %qkw9_1_bi = index.constant 156 : index + %qkw9_1_b = vector.load %qkt_flat[%qkw9_1_bi] : view<256xf32> -> vector<4xf32> + %qkd9_9 = vector.extract %qkw9_1_a[1] : vector<4xf32> -> f32 + %qsum9 = scalar.fmaf %qkd9_9, %dl9, %aq8_9 : f32 + %ao9_r = scalar.fmaf %bvv9, %cp9, %qsum9 : f32 + %ao9 = scf.if %active9 -> (f32) { + %value = scalar.mulf %ao9_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi9_a = index.add %rec_col, %k1152 : index + %aoi9 = index.assume %aoi9_a [range(%aoi9_a, 0, 2047)] : index + %ddh9 = scalar.fptrunc %dh9 : f32 to f16 + %publish_active9 = scalar.andi %lane_live, %active9 : i1 + scf.if %publish_active9 { + view.store %ao9, %aost_flat[%aoi9] : f32, view<2048xf32> + } + %gp9_10 = vector.extract %grw9_1_a[2] : vector<4xf32> -> f32 + %ag9_10 = scalar.fmaf %gp9_10, %dl9, %ag8_10 : f32 + %qp9_10 = vector.extract %qkw9_1_a[2] : vector<4xf32> -> f32 + %aq9_10 = scalar.fmaf %qp9_10, %dl9, %aq8_10 : f32 + %gp9_11 = vector.extract %grw9_1_a[3] : vector<4xf32> -> f32 + %ag9_11 = scalar.fmaf %gp9_11, %dl9, %ag8_11 : f32 + %qp9_11 = vector.extract %qkw9_1_a[3] : vector<4xf32> -> f32 + %aq9_11 = scalar.fmaf %qp9_11, %dl9, %aq8_11 : f32 + %gp9_12 = vector.extract %grw9_1_b[0] : vector<4xf32> -> f32 + %ag9_12 = scalar.fmaf %gp9_12, %dl9, %ag8_12 : f32 + %qp9_12 = vector.extract %qkw9_1_b[0] : vector<4xf32> -> f32 + %aq9_12 = scalar.fmaf %qp9_12, %dl9, %aq8_12 : f32 + %gp9_13 = vector.extract %grw9_1_b[1] : vector<4xf32> -> f32 + %ag9_13 = scalar.fmaf %gp9_13, %dl9, %ag8_13 : f32 + %qp9_13 = vector.extract %qkw9_1_b[1] : vector<4xf32> -> f32 + %aq9_13 = scalar.fmaf %qp9_13, %dl9, %aq8_13 : f32 + %gp9_14 = vector.extract %grw9_1_b[2] : vector<4xf32> -> f32 + %ag9_14 = scalar.fmaf %gp9_14, %dl9, %ag8_14 : f32 + %qp9_14 = vector.extract %qkw9_1_b[2] : vector<4xf32> -> f32 + %aq9_14 = scalar.fmaf %qp9_14, %dl9, %aq8_14 : f32 + %gp9_15 = vector.extract %grw9_1_b[3] : vector<4xf32> -> f32 + %ag9_15 = scalar.fmaf %gp9_15, %dl9, %ag8_15 : f32 + %qp9_15 = vector.extract %qkw9_1_b[3] : vector<4xf32> -> f32 + %aq9_15 = scalar.fmaf %qp9_15, %dl9, %aq8_15 : f32 + %gq10 = vector.load %gs_flat[%k40] : view<80xf32> -> vector<4xf32> + %bv10 = vector.extract %gq10[0] : vector<4xf32> -> f32 + %cp10 = vector.extract %gq10[1] : vector<4xf32> -> f32 + %ev10 = vector.extract %gq10[2] : vector<4xf32> -> f32 + %tk10_r = index.add %t0, %k10 : index + %tk10_ok = index.cmp ult, %tk10_r, %n_tokens : index + %tk10_s = scf.select %tk10_ok, %tk10_r, %c0 : index + %tk10 = index.assume %tk10_s [range(%tk10_s, 0, 1048575)] : index + %av10_i = index.add %tile_row, %k10 : index + %av10_b = index.assume %av10_i [range(%av10_i, 0, 271)] : index + %av10 = view.load %a_flat[%av10_b] : view<272xf32> -> f32 + %bvv10 = view.load %b_flat[%av10_b] : view<272xf32> -> f32 + %vrd10_a = index.add %rec_col, %k1280 : index + %vrd10 = index.assume %vrd10_a [range(%vrd10_a, 0, 2047)] : index + %vval10 = view.load %vst_flat[%vrd10] : view<2048xf32> -> f32 + %adec10 = scalar.fmaf %av10, %cp10, %ag9_10 : f32 + %vsub10 = scalar.subf %vval10, %adec10 : f32 + %active_index10 = index.add %t0, %k10 : index + %row_in_range10 = index.cmp ult, %active_index10, %n_tokens : index + %active10 = scalar.ori %row_in_range10, %all_recurrence_rows : i1 + %dl10 = scf.if %active10 -> (f32) { + %value = scalar.mulf %vsub10, %bv10 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval10, %bv10 : f32 + scf.yield %padded : f32 + } + %dh10 = scalar.mulf %dl10, %ev10 : f32 + %grw10_0_a = vector.load %gr_flat[%k160] : view<256xf32> -> vector<4xf32> + %grw10_0_bi = index.constant 164 : index + %grw10_0_b = vector.load %gr_flat[%grw10_0_bi] : view<256xf32> -> vector<4xf32> + %grw10_1_a = vector.load %gr_flat[%k168] : view<256xf32> -> vector<4xf32> + %grw10_1_bi = index.constant 172 : index + %grw10_1_b = vector.load %gr_flat[%grw10_1_bi] : view<256xf32> -> vector<4xf32> + %qkw10_0_a = vector.load %qkt_flat[%k160] : view<256xf32> -> vector<4xf32> + %qkw10_0_bi = index.constant 164 : index + %qkw10_0_b = vector.load %qkt_flat[%qkw10_0_bi] : view<256xf32> -> vector<4xf32> + %qkw10_1_a = vector.load %qkt_flat[%k168] : view<256xf32> -> vector<4xf32> + %qkw10_1_bi = index.constant 172 : index + %qkw10_1_b = vector.load %qkt_flat[%qkw10_1_bi] : view<256xf32> -> vector<4xf32> + %qkd10_10 = vector.extract %qkw10_1_a[2] : vector<4xf32> -> f32 + %qsum10 = scalar.fmaf %qkd10_10, %dl10, %aq9_10 : f32 + %ao10_r = scalar.fmaf %bvv10, %cp10, %qsum10 : f32 + %ao10 = scf.if %active10 -> (f32) { + %value = scalar.mulf %ao10_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi10_a = index.add %rec_col, %k1280 : index + %aoi10 = index.assume %aoi10_a [range(%aoi10_a, 0, 2047)] : index + %ddh10 = scalar.fptrunc %dh10 : f32 to f16 + %publish_active10 = scalar.andi %lane_live, %active10 : i1 + scf.if %publish_active10 { + view.store %ao10, %aost_flat[%aoi10] : f32, view<2048xf32> + } + %gp10_11 = vector.extract %grw10_1_a[3] : vector<4xf32> -> f32 + %ag10_11 = scalar.fmaf %gp10_11, %dl10, %ag9_11 : f32 + %qp10_11 = vector.extract %qkw10_1_a[3] : vector<4xf32> -> f32 + %aq10_11 = scalar.fmaf %qp10_11, %dl10, %aq9_11 : f32 + %gp10_12 = vector.extract %grw10_1_b[0] : vector<4xf32> -> f32 + %ag10_12 = scalar.fmaf %gp10_12, %dl10, %ag9_12 : f32 + %qp10_12 = vector.extract %qkw10_1_b[0] : vector<4xf32> -> f32 + %aq10_12 = scalar.fmaf %qp10_12, %dl10, %aq9_12 : f32 + %gp10_13 = vector.extract %grw10_1_b[1] : vector<4xf32> -> f32 + %ag10_13 = scalar.fmaf %gp10_13, %dl10, %ag9_13 : f32 + %qp10_13 = vector.extract %qkw10_1_b[1] : vector<4xf32> -> f32 + %aq10_13 = scalar.fmaf %qp10_13, %dl10, %aq9_13 : f32 + %gp10_14 = vector.extract %grw10_1_b[2] : vector<4xf32> -> f32 + %ag10_14 = scalar.fmaf %gp10_14, %dl10, %ag9_14 : f32 + %qp10_14 = vector.extract %qkw10_1_b[2] : vector<4xf32> -> f32 + %aq10_14 = scalar.fmaf %qp10_14, %dl10, %aq9_14 : f32 + %gp10_15 = vector.extract %grw10_1_b[3] : vector<4xf32> -> f32 + %ag10_15 = scalar.fmaf %gp10_15, %dl10, %ag9_15 : f32 + %qp10_15 = vector.extract %qkw10_1_b[3] : vector<4xf32> -> f32 + %aq10_15 = scalar.fmaf %qp10_15, %dl10, %aq9_15 : f32 + %gq11 = vector.load %gs_flat[%k44] : view<80xf32> -> vector<4xf32> + %bv11 = vector.extract %gq11[0] : vector<4xf32> -> f32 + %cp11 = vector.extract %gq11[1] : vector<4xf32> -> f32 + %ev11 = vector.extract %gq11[2] : vector<4xf32> -> f32 + %tk11_r = index.add %t0, %k11 : index + %tk11_ok = index.cmp ult, %tk11_r, %n_tokens : index + %tk11_s = scf.select %tk11_ok, %tk11_r, %c0 : index + %tk11 = index.assume %tk11_s [range(%tk11_s, 0, 1048575)] : index + %av11_i = index.add %tile_row, %k11 : index + %av11_b = index.assume %av11_i [range(%av11_i, 0, 271)] : index + %av11 = view.load %a_flat[%av11_b] : view<272xf32> -> f32 + %bvv11 = view.load %b_flat[%av11_b] : view<272xf32> -> f32 + %vrd11_a = index.add %rec_col, %k1408 : index + %vrd11 = index.assume %vrd11_a [range(%vrd11_a, 0, 2047)] : index + %vval11 = view.load %vst_flat[%vrd11] : view<2048xf32> -> f32 + %adec11 = scalar.fmaf %av11, %cp11, %ag10_11 : f32 + %vsub11 = scalar.subf %vval11, %adec11 : f32 + %active_index11 = index.add %t0, %k11 : index + %row_in_range11 = index.cmp ult, %active_index11, %n_tokens : index + %active11 = scalar.ori %row_in_range11, %all_recurrence_rows : i1 + %dl11 = scf.if %active11 -> (f32) { + %value = scalar.mulf %vsub11, %bv11 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval11, %bv11 : f32 + scf.yield %padded : f32 + } + %dh11 = scalar.mulf %dl11, %ev11 : f32 + %grw11_0_a = vector.load %gr_flat[%k176] : view<256xf32> -> vector<4xf32> + %grw11_0_bi = index.constant 180 : index + %grw11_0_b = vector.load %gr_flat[%grw11_0_bi] : view<256xf32> -> vector<4xf32> + %grw11_1_a = vector.load %gr_flat[%k184] : view<256xf32> -> vector<4xf32> + %grw11_1_bi = index.constant 188 : index + %grw11_1_b = vector.load %gr_flat[%grw11_1_bi] : view<256xf32> -> vector<4xf32> + %qkw11_0_a = vector.load %qkt_flat[%k176] : view<256xf32> -> vector<4xf32> + %qkw11_0_bi = index.constant 180 : index + %qkw11_0_b = vector.load %qkt_flat[%qkw11_0_bi] : view<256xf32> -> vector<4xf32> + %qkw11_1_a = vector.load %qkt_flat[%k184] : view<256xf32> -> vector<4xf32> + %qkw11_1_bi = index.constant 188 : index + %qkw11_1_b = vector.load %qkt_flat[%qkw11_1_bi] : view<256xf32> -> vector<4xf32> + %qkd11_11 = vector.extract %qkw11_1_a[3] : vector<4xf32> -> f32 + %qsum11 = scalar.fmaf %qkd11_11, %dl11, %aq10_11 : f32 + %ao11_r = scalar.fmaf %bvv11, %cp11, %qsum11 : f32 + %ao11 = scf.if %active11 -> (f32) { + %value = scalar.mulf %ao11_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi11_a = index.add %rec_col, %k1408 : index + %aoi11 = index.assume %aoi11_a [range(%aoi11_a, 0, 2047)] : index + %ddh11 = scalar.fptrunc %dh11 : f32 to f16 + %publish_active11 = scalar.andi %lane_live, %active11 : i1 + scf.if %publish_active11 { + view.store %ao11, %aost_flat[%aoi11] : f32, view<2048xf32> + } + %gp11_12 = vector.extract %grw11_1_b[0] : vector<4xf32> -> f32 + %ag11_12 = scalar.fmaf %gp11_12, %dl11, %ag10_12 : f32 + %qp11_12 = vector.extract %qkw11_1_b[0] : vector<4xf32> -> f32 + %aq11_12 = scalar.fmaf %qp11_12, %dl11, %aq10_12 : f32 + %gp11_13 = vector.extract %grw11_1_b[1] : vector<4xf32> -> f32 + %ag11_13 = scalar.fmaf %gp11_13, %dl11, %ag10_13 : f32 + %qp11_13 = vector.extract %qkw11_1_b[1] : vector<4xf32> -> f32 + %aq11_13 = scalar.fmaf %qp11_13, %dl11, %aq10_13 : f32 + %gp11_14 = vector.extract %grw11_1_b[2] : vector<4xf32> -> f32 + %ag11_14 = scalar.fmaf %gp11_14, %dl11, %ag10_14 : f32 + %qp11_14 = vector.extract %qkw11_1_b[2] : vector<4xf32> -> f32 + %aq11_14 = scalar.fmaf %qp11_14, %dl11, %aq10_14 : f32 + %gp11_15 = vector.extract %grw11_1_b[3] : vector<4xf32> -> f32 + %ag11_15 = scalar.fmaf %gp11_15, %dl11, %ag10_15 : f32 + %qp11_15 = vector.extract %qkw11_1_b[3] : vector<4xf32> -> f32 + %aq11_15 = scalar.fmaf %qp11_15, %dl11, %aq10_15 : f32 + %gq12 = vector.load %gs_flat[%k48] : view<80xf32> -> vector<4xf32> + %bv12 = vector.extract %gq12[0] : vector<4xf32> -> f32 + %cp12 = vector.extract %gq12[1] : vector<4xf32> -> f32 + %ev12 = vector.extract %gq12[2] : vector<4xf32> -> f32 + %tk12_r = index.add %t0, %k12 : index + %tk12_ok = index.cmp ult, %tk12_r, %n_tokens : index + %tk12_s = scf.select %tk12_ok, %tk12_r, %c0 : index + %tk12 = index.assume %tk12_s [range(%tk12_s, 0, 1048575)] : index + %av12_i = index.add %tile_row, %k12 : index + %av12_b = index.assume %av12_i [range(%av12_i, 0, 271)] : index + %av12 = view.load %a_flat[%av12_b] : view<272xf32> -> f32 + %bvv12 = view.load %b_flat[%av12_b] : view<272xf32> -> f32 + %vrd12_a = index.add %rec_col, %k1536 : index + %vrd12 = index.assume %vrd12_a [range(%vrd12_a, 0, 2047)] : index + %vval12 = view.load %vst_flat[%vrd12] : view<2048xf32> -> f32 + %adec12 = scalar.fmaf %av12, %cp12, %ag11_12 : f32 + %vsub12 = scalar.subf %vval12, %adec12 : f32 + %active_index12 = index.add %t0, %k12 : index + %row_in_range12 = index.cmp ult, %active_index12, %n_tokens : index + %active12 = scalar.ori %row_in_range12, %all_recurrence_rows : i1 + %dl12 = scf.if %active12 -> (f32) { + %value = scalar.mulf %vsub12, %bv12 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval12, %bv12 : f32 + scf.yield %padded : f32 + } + %dh12 = scalar.mulf %dl12, %ev12 : f32 + %grw12_0_a = vector.load %gr_flat[%k192] : view<256xf32> -> vector<4xf32> + %grw12_0_bi = index.constant 196 : index + %grw12_0_b = vector.load %gr_flat[%grw12_0_bi] : view<256xf32> -> vector<4xf32> + %grw12_1_a = vector.load %gr_flat[%k200] : view<256xf32> -> vector<4xf32> + %grw12_1_bi = index.constant 204 : index + %grw12_1_b = vector.load %gr_flat[%grw12_1_bi] : view<256xf32> -> vector<4xf32> + %qkw12_0_a = vector.load %qkt_flat[%k192] : view<256xf32> -> vector<4xf32> + %qkw12_0_bi = index.constant 196 : index + %qkw12_0_b = vector.load %qkt_flat[%qkw12_0_bi] : view<256xf32> -> vector<4xf32> + %qkw12_1_a = vector.load %qkt_flat[%k200] : view<256xf32> -> vector<4xf32> + %qkw12_1_bi = index.constant 204 : index + %qkw12_1_b = vector.load %qkt_flat[%qkw12_1_bi] : view<256xf32> -> vector<4xf32> + %qkd12_12 = vector.extract %qkw12_1_b[0] : vector<4xf32> -> f32 + %qsum12 = scalar.fmaf %qkd12_12, %dl12, %aq11_12 : f32 + %ao12_r = scalar.fmaf %bvv12, %cp12, %qsum12 : f32 + %ao12 = scf.if %active12 -> (f32) { + %value = scalar.mulf %ao12_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi12_a = index.add %rec_col, %k1536 : index + %aoi12 = index.assume %aoi12_a [range(%aoi12_a, 0, 2047)] : index + %ddh12 = scalar.fptrunc %dh12 : f32 to f16 + %publish_active12 = scalar.andi %lane_live, %active12 : i1 + scf.if %publish_active12 { + view.store %ao12, %aost_flat[%aoi12] : f32, view<2048xf32> + } + %gp12_13 = vector.extract %grw12_1_b[1] : vector<4xf32> -> f32 + %ag12_13 = scalar.fmaf %gp12_13, %dl12, %ag11_13 : f32 + %qp12_13 = vector.extract %qkw12_1_b[1] : vector<4xf32> -> f32 + %aq12_13 = scalar.fmaf %qp12_13, %dl12, %aq11_13 : f32 + %gp12_14 = vector.extract %grw12_1_b[2] : vector<4xf32> -> f32 + %ag12_14 = scalar.fmaf %gp12_14, %dl12, %ag11_14 : f32 + %qp12_14 = vector.extract %qkw12_1_b[2] : vector<4xf32> -> f32 + %aq12_14 = scalar.fmaf %qp12_14, %dl12, %aq11_14 : f32 + %gp12_15 = vector.extract %grw12_1_b[3] : vector<4xf32> -> f32 + %ag12_15 = scalar.fmaf %gp12_15, %dl12, %ag11_15 : f32 + %qp12_15 = vector.extract %qkw12_1_b[3] : vector<4xf32> -> f32 + %aq12_15 = scalar.fmaf %qp12_15, %dl12, %aq11_15 : f32 + %gq13 = vector.load %gs_flat[%k52] : view<80xf32> -> vector<4xf32> + %bv13 = vector.extract %gq13[0] : vector<4xf32> -> f32 + %cp13 = vector.extract %gq13[1] : vector<4xf32> -> f32 + %ev13 = vector.extract %gq13[2] : vector<4xf32> -> f32 + %tk13_r = index.add %t0, %k13 : index + %tk13_ok = index.cmp ult, %tk13_r, %n_tokens : index + %tk13_s = scf.select %tk13_ok, %tk13_r, %c0 : index + %tk13 = index.assume %tk13_s [range(%tk13_s, 0, 1048575)] : index + %av13_i = index.add %tile_row, %k13 : index + %av13_b = index.assume %av13_i [range(%av13_i, 0, 271)] : index + %av13 = view.load %a_flat[%av13_b] : view<272xf32> -> f32 + %bvv13 = view.load %b_flat[%av13_b] : view<272xf32> -> f32 + %vrd13_a = index.add %rec_col, %k1664 : index + %vrd13 = index.assume %vrd13_a [range(%vrd13_a, 0, 2047)] : index + %vval13 = view.load %vst_flat[%vrd13] : view<2048xf32> -> f32 + %adec13 = scalar.fmaf %av13, %cp13, %ag12_13 : f32 + %vsub13 = scalar.subf %vval13, %adec13 : f32 + %active_index13 = index.add %t0, %k13 : index + %row_in_range13 = index.cmp ult, %active_index13, %n_tokens : index + %active13 = scalar.ori %row_in_range13, %all_recurrence_rows : i1 + %dl13 = scf.if %active13 -> (f32) { + %value = scalar.mulf %vsub13, %bv13 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval13, %bv13 : f32 + scf.yield %padded : f32 + } + %dh13 = scalar.mulf %dl13, %ev13 : f32 + %grw13_0_a = vector.load %gr_flat[%k208] : view<256xf32> -> vector<4xf32> + %grw13_0_bi = index.constant 212 : index + %grw13_0_b = vector.load %gr_flat[%grw13_0_bi] : view<256xf32> -> vector<4xf32> + %grw13_1_a = vector.load %gr_flat[%k216] : view<256xf32> -> vector<4xf32> + %grw13_1_bi = index.constant 220 : index + %grw13_1_b = vector.load %gr_flat[%grw13_1_bi] : view<256xf32> -> vector<4xf32> + %qkw13_0_a = vector.load %qkt_flat[%k208] : view<256xf32> -> vector<4xf32> + %qkw13_0_bi = index.constant 212 : index + %qkw13_0_b = vector.load %qkt_flat[%qkw13_0_bi] : view<256xf32> -> vector<4xf32> + %qkw13_1_a = vector.load %qkt_flat[%k216] : view<256xf32> -> vector<4xf32> + %qkw13_1_bi = index.constant 220 : index + %qkw13_1_b = vector.load %qkt_flat[%qkw13_1_bi] : view<256xf32> -> vector<4xf32> + %qkd13_13 = vector.extract %qkw13_1_b[1] : vector<4xf32> -> f32 + %qsum13 = scalar.fmaf %qkd13_13, %dl13, %aq12_13 : f32 + %ao13_r = scalar.fmaf %bvv13, %cp13, %qsum13 : f32 + %ao13 = scf.if %active13 -> (f32) { + %value = scalar.mulf %ao13_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi13_a = index.add %rec_col, %k1664 : index + %aoi13 = index.assume %aoi13_a [range(%aoi13_a, 0, 2047)] : index + %ddh13 = scalar.fptrunc %dh13 : f32 to f16 + %publish_active13 = scalar.andi %lane_live, %active13 : i1 + scf.if %publish_active13 { + view.store %ao13, %aost_flat[%aoi13] : f32, view<2048xf32> + } + %gp13_14 = vector.extract %grw13_1_b[2] : vector<4xf32> -> f32 + %ag13_14 = scalar.fmaf %gp13_14, %dl13, %ag12_14 : f32 + %qp13_14 = vector.extract %qkw13_1_b[2] : vector<4xf32> -> f32 + %aq13_14 = scalar.fmaf %qp13_14, %dl13, %aq12_14 : f32 + %gp13_15 = vector.extract %grw13_1_b[3] : vector<4xf32> -> f32 + %ag13_15 = scalar.fmaf %gp13_15, %dl13, %ag12_15 : f32 + %qp13_15 = vector.extract %qkw13_1_b[3] : vector<4xf32> -> f32 + %aq13_15 = scalar.fmaf %qp13_15, %dl13, %aq12_15 : f32 + %gq14 = vector.load %gs_flat[%k56] : view<80xf32> -> vector<4xf32> + %bv14 = vector.extract %gq14[0] : vector<4xf32> -> f32 + %cp14 = vector.extract %gq14[1] : vector<4xf32> -> f32 + %ev14 = vector.extract %gq14[2] : vector<4xf32> -> f32 + %tk14_r = index.add %t0, %k14 : index + %tk14_ok = index.cmp ult, %tk14_r, %n_tokens : index + %tk14_s = scf.select %tk14_ok, %tk14_r, %c0 : index + %tk14 = index.assume %tk14_s [range(%tk14_s, 0, 1048575)] : index + %av14_i = index.add %tile_row, %k14 : index + %av14_b = index.assume %av14_i [range(%av14_i, 0, 271)] : index + %av14 = view.load %a_flat[%av14_b] : view<272xf32> -> f32 + %bvv14 = view.load %b_flat[%av14_b] : view<272xf32> -> f32 + %vrd14_a = index.add %rec_col, %k1792 : index + %vrd14 = index.assume %vrd14_a [range(%vrd14_a, 0, 2047)] : index + %vval14 = view.load %vst_flat[%vrd14] : view<2048xf32> -> f32 + %adec14 = scalar.fmaf %av14, %cp14, %ag13_14 : f32 + %vsub14 = scalar.subf %vval14, %adec14 : f32 + %active_index14 = index.add %t0, %k14 : index + %row_in_range14 = index.cmp ult, %active_index14, %n_tokens : index + %active14 = scalar.ori %row_in_range14, %all_recurrence_rows : i1 + %dl14 = scf.if %active14 -> (f32) { + %value = scalar.mulf %vsub14, %bv14 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval14, %bv14 : f32 + scf.yield %padded : f32 + } + %dh14 = scalar.mulf %dl14, %ev14 : f32 + %grw14_0_a = vector.load %gr_flat[%k224] : view<256xf32> -> vector<4xf32> + %grw14_0_bi = index.constant 228 : index + %grw14_0_b = vector.load %gr_flat[%grw14_0_bi] : view<256xf32> -> vector<4xf32> + %grw14_1_a = vector.load %gr_flat[%k232] : view<256xf32> -> vector<4xf32> + %grw14_1_bi = index.constant 236 : index + %grw14_1_b = vector.load %gr_flat[%grw14_1_bi] : view<256xf32> -> vector<4xf32> + %qkw14_0_a = vector.load %qkt_flat[%k224] : view<256xf32> -> vector<4xf32> + %qkw14_0_bi = index.constant 228 : index + %qkw14_0_b = vector.load %qkt_flat[%qkw14_0_bi] : view<256xf32> -> vector<4xf32> + %qkw14_1_a = vector.load %qkt_flat[%k232] : view<256xf32> -> vector<4xf32> + %qkw14_1_bi = index.constant 236 : index + %qkw14_1_b = vector.load %qkt_flat[%qkw14_1_bi] : view<256xf32> -> vector<4xf32> + %qkd14_14 = vector.extract %qkw14_1_b[2] : vector<4xf32> -> f32 + %qsum14 = scalar.fmaf %qkd14_14, %dl14, %aq13_14 : f32 + %ao14_r = scalar.fmaf %bvv14, %cp14, %qsum14 : f32 + %ao14 = scf.if %active14 -> (f32) { + %value = scalar.mulf %ao14_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi14_a = index.add %rec_col, %k1792 : index + %aoi14 = index.assume %aoi14_a [range(%aoi14_a, 0, 2047)] : index + %ddh14 = scalar.fptrunc %dh14 : f32 to f16 + %publish_active14 = scalar.andi %lane_live, %active14 : i1 + scf.if %publish_active14 { + view.store %ao14, %aost_flat[%aoi14] : f32, view<2048xf32> + } + %gp14_15 = vector.extract %grw14_1_b[3] : vector<4xf32> -> f32 + %ag14_15 = scalar.fmaf %gp14_15, %dl14, %ag13_15 : f32 + %qp14_15 = vector.extract %qkw14_1_b[3] : vector<4xf32> -> f32 + %aq14_15 = scalar.fmaf %qp14_15, %dl14, %aq13_15 : f32 + %gq15 = vector.load %gs_flat[%k60] : view<80xf32> -> vector<4xf32> + %bv15 = vector.extract %gq15[0] : vector<4xf32> -> f32 + %cp15 = vector.extract %gq15[1] : vector<4xf32> -> f32 + %ev15 = vector.extract %gq15[2] : vector<4xf32> -> f32 + %tk15_r = index.add %t0, %k15 : index + %tk15_ok = index.cmp ult, %tk15_r, %n_tokens : index + %tk15_s = scf.select %tk15_ok, %tk15_r, %c0 : index + %tk15 = index.assume %tk15_s [range(%tk15_s, 0, 1048575)] : index + %av15_i = index.add %tile_row, %k15 : index + %av15_b = index.assume %av15_i [range(%av15_i, 0, 271)] : index + %av15 = view.load %a_flat[%av15_b] : view<272xf32> -> f32 + %bvv15 = view.load %b_flat[%av15_b] : view<272xf32> -> f32 + %vrd15_a = index.add %rec_col, %k1920 : index + %vrd15 = index.assume %vrd15_a [range(%vrd15_a, 0, 2047)] : index + %vval15 = view.load %vst_flat[%vrd15] : view<2048xf32> -> f32 + %adec15 = scalar.fmaf %av15, %cp15, %ag14_15 : f32 + %vsub15 = scalar.subf %vval15, %adec15 : f32 + %active_index15 = index.add %t0, %k15 : index + %row_in_range15 = index.cmp ult, %active_index15, %n_tokens : index + %active15 = scalar.ori %row_in_range15, %all_recurrence_rows : i1 + %dl15 = scf.if %active15 -> (f32) { + %value = scalar.mulf %vsub15, %bv15 : f32 + scf.yield %value : f32 + } else { + %padded = scalar.mulf %vval15, %bv15 : f32 + scf.yield %padded : f32 + } + %dh15 = scalar.mulf %dl15, %ev15 : f32 + %grw15_0_a = vector.load %gr_flat[%k240] : view<256xf32> -> vector<4xf32> + %grw15_0_bi = index.constant 244 : index + %grw15_0_b = vector.load %gr_flat[%grw15_0_bi] : view<256xf32> -> vector<4xf32> + %grw15_1_a = vector.load %gr_flat[%k248] : view<256xf32> -> vector<4xf32> + %grw15_1_bi = index.constant 252 : index + %grw15_1_b = vector.load %gr_flat[%grw15_1_bi] : view<256xf32> -> vector<4xf32> + %qkw15_0_a = vector.load %qkt_flat[%k240] : view<256xf32> -> vector<4xf32> + %qkw15_0_bi = index.constant 244 : index + %qkw15_0_b = vector.load %qkt_flat[%qkw15_0_bi] : view<256xf32> -> vector<4xf32> + %qkw15_1_a = vector.load %qkt_flat[%k248] : view<256xf32> -> vector<4xf32> + %qkw15_1_bi = index.constant 252 : index + %qkw15_1_b = vector.load %qkt_flat[%qkw15_1_bi] : view<256xf32> -> vector<4xf32> + %qkd15_15 = vector.extract %qkw15_1_b[3] : vector<4xf32> -> f32 + %qsum15 = scalar.fmaf %qkd15_15, %dl15, %aq14_15 : f32 + %ao15_r = scalar.fmaf %bvv15, %cp15, %qsum15 : f32 + %ao15 = scf.if %active15 -> (f32) { + %value = scalar.mulf %ao15_r, %scale : f32 + scf.yield %value : f32 + } else { + %zero = scalar.constant 0.0 : f32 + scf.yield %zero : f32 + } + %aoi15_a = index.add %rec_col, %k1920 : index + %aoi15 = index.assume %aoi15_a [range(%aoi15_a, 0, 2047)] : index + %ddh15 = scalar.fptrunc %dh15 : f32 to f16 + %publish_active15 = scalar.andi %lane_live, %active15 : i1 + scf.if %publish_active15 { + view.store %ao15, %aost_flat[%aoi15] : f32, view<2048xf32> + } + // A is now dead. Materialize the register-held deltas into its + // allocation so the full-head owner does not reserve a separate D tile. + %ddv0 = vector.from_elements %ddh0, %ddh1, %ddh2, %ddh3, %ddh4, %ddh5, %ddh6, %ddh7 : vector<8xf16> + %ddv1 = vector.from_elements %ddh8, %ddh9, %ddh10, %ddh11, %ddh12, %ddh13, %ddh14, %ddh15 : vector<8xf16> + %dds0 = index.assume %d_row [range(%d_row, 0, 376)] : index + %dds1_a = index.add %d_row, %k8 : index + %dds1 = index.assume %dds1_a [range(%dds1_a, 8, 376)] : index + scf.if %lane_live { + vector.store %ddv0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %ddv1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + scf.if %rmsnorm_gate { + scf.if %publish_q8 { + } else { + template.apply<@llm.gated_delta_net.rmsnorm_gate.publish_head128>(%base, %n_tokens, %n_heads, %n_seqs, %lane, %wave, %t0, %seq, %c0, %h_idx, %k4, %k8, %k16, %k128, %rms_epsilon, %rms_gate_op, %publish_q8, %lds, %rms_weight, %raw_gate, %norm_output, %half_output) : (offset, index, index, index, index, index, index, index, index, index, index, index, index, index, f32, index, i1, buffer, buffer, buffer, buffer, buffer) + + } + } else { + %aow0_r = index.add %t0, %k0 : index + %aow0_s = index.add %aow0_r, %vs_tok : index + %aow0_ok = index.cmp ult, %aow0_s, %n_tokens : index + %aow0_c = scf.select %aow0_ok, %aow0_s, %c0 : index + %aow0 = index.assume %aow0_c [range(%aow0_c, 0, 1048575)] : index + %aol0_m = index.mul %vs_tok, %k128 : index + %aol0_a = index.add %aol0_m, %vs_col : index + %aol0_b = index.add %aol0_a, %k0 : index + %aol0 = index.assume %aol0_b [range(%aol0_b, 0, 2047)] : index + %aov0 = view.load %aost_flat[%aol0] : view<2048xf32> -> f32 + %aog0_m = index.mul %aow0, %attn_stride : index + %aog0_a = index.add %attn_base, %aog0_m : index + %aog0_c = index.add %aog0_a, %vs_gcol : index + %aog0 = index.assume %aog0_c [range(%aog0_c, 0, 1073741823)] : index + scf.if %aow0_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov0, %dst_view, %aog0) : (f32, view<1073741824xf32>, index) + } + %aow1_r = index.add %t0, %k2 : index + %aow1_s = index.add %aow1_r, %vs_tok : index + %aow1_ok = index.cmp ult, %aow1_s, %n_tokens : index + %aow1_c = scf.select %aow1_ok, %aow1_s, %c0 : index + %aow1 = index.assume %aow1_c [range(%aow1_c, 0, 1048575)] : index + %aol1_m = index.mul %vs_tok, %k128 : index + %aol1_a = index.add %aol1_m, %vs_col : index + %aol1_b = index.add %aol1_a, %k256 : index + %aol1 = index.assume %aol1_b [range(%aol1_b, 0, 2047)] : index + %aov1 = view.load %aost_flat[%aol1] : view<2048xf32> -> f32 + %aog1_m = index.mul %aow1, %attn_stride : index + %aog1_a = index.add %attn_base, %aog1_m : index + %aog1_c = index.add %aog1_a, %vs_gcol : index + %aog1 = index.assume %aog1_c [range(%aog1_c, 0, 1073741823)] : index + scf.if %aow1_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov1, %dst_view, %aog1) : (f32, view<1073741824xf32>, index) + } + %aow2_r = index.add %t0, %k4 : index + %aow2_s = index.add %aow2_r, %vs_tok : index + %aow2_ok = index.cmp ult, %aow2_s, %n_tokens : index + %aow2_c = scf.select %aow2_ok, %aow2_s, %c0 : index + %aow2 = index.assume %aow2_c [range(%aow2_c, 0, 1048575)] : index + %aol2_m = index.mul %vs_tok, %k128 : index + %aol2_a = index.add %aol2_m, %vs_col : index + %aol2_b = index.add %aol2_a, %k512 : index + %aol2 = index.assume %aol2_b [range(%aol2_b, 0, 2047)] : index + %aov2 = view.load %aost_flat[%aol2] : view<2048xf32> -> f32 + %aog2_m = index.mul %aow2, %attn_stride : index + %aog2_a = index.add %attn_base, %aog2_m : index + %aog2_c = index.add %aog2_a, %vs_gcol : index + %aog2 = index.assume %aog2_c [range(%aog2_c, 0, 1073741823)] : index + scf.if %aow2_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov2, %dst_view, %aog2) : (f32, view<1073741824xf32>, index) + } + %aow3_r = index.add %t0, %k6 : index + %aow3_s = index.add %aow3_r, %vs_tok : index + %aow3_ok = index.cmp ult, %aow3_s, %n_tokens : index + %aow3_c = scf.select %aow3_ok, %aow3_s, %c0 : index + %aow3 = index.assume %aow3_c [range(%aow3_c, 0, 1048575)] : index + %aol3_m = index.mul %vs_tok, %k128 : index + %aol3_a = index.add %aol3_m, %vs_col : index + %aol3_b = index.add %aol3_a, %k768 : index + %aol3 = index.assume %aol3_b [range(%aol3_b, 0, 2047)] : index + %aov3 = view.load %aost_flat[%aol3] : view<2048xf32> -> f32 + %aog3_m = index.mul %aow3, %attn_stride : index + %aog3_a = index.add %attn_base, %aog3_m : index + %aog3_c = index.add %aog3_a, %vs_gcol : index + %aog3 = index.assume %aog3_c [range(%aog3_c, 0, 1073741823)] : index + scf.if %aow3_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov3, %dst_view, %aog3) : (f32, view<1073741824xf32>, index) + } + %aow4_r = index.add %t0, %k8 : index + %aow4_s = index.add %aow4_r, %vs_tok : index + %aow4_ok = index.cmp ult, %aow4_s, %n_tokens : index + %aow4_c = scf.select %aow4_ok, %aow4_s, %c0 : index + %aow4 = index.assume %aow4_c [range(%aow4_c, 0, 1048575)] : index + %aol4_m = index.mul %vs_tok, %k128 : index + %aol4_a = index.add %aol4_m, %vs_col : index + %aol4_b = index.add %aol4_a, %k1024 : index + %aol4 = index.assume %aol4_b [range(%aol4_b, 0, 2047)] : index + %aov4 = view.load %aost_flat[%aol4] : view<2048xf32> -> f32 + %aog4_m = index.mul %aow4, %attn_stride : index + %aog4_a = index.add %attn_base, %aog4_m : index + %aog4_c = index.add %aog4_a, %vs_gcol : index + %aog4 = index.assume %aog4_c [range(%aog4_c, 0, 1073741823)] : index + scf.if %aow4_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov4, %dst_view, %aog4) : (f32, view<1073741824xf32>, index) + } + %aow5_r = index.add %t0, %k10 : index + %aow5_s = index.add %aow5_r, %vs_tok : index + %aow5_ok = index.cmp ult, %aow5_s, %n_tokens : index + %aow5_c = scf.select %aow5_ok, %aow5_s, %c0 : index + %aow5 = index.assume %aow5_c [range(%aow5_c, 0, 1048575)] : index + %aol5_m = index.mul %vs_tok, %k128 : index + %aol5_a = index.add %aol5_m, %vs_col : index + %aol5_b = index.add %aol5_a, %k1280 : index + %aol5 = index.assume %aol5_b [range(%aol5_b, 0, 2047)] : index + %aov5 = view.load %aost_flat[%aol5] : view<2048xf32> -> f32 + %aog5_m = index.mul %aow5, %attn_stride : index + %aog5_a = index.add %attn_base, %aog5_m : index + %aog5_c = index.add %aog5_a, %vs_gcol : index + %aog5 = index.assume %aog5_c [range(%aog5_c, 0, 1073741823)] : index + scf.if %aow5_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov5, %dst_view, %aog5) : (f32, view<1073741824xf32>, index) + } + %aow6_r = index.add %t0, %k12 : index + %aow6_s = index.add %aow6_r, %vs_tok : index + %aow6_ok = index.cmp ult, %aow6_s, %n_tokens : index + %aow6_c = scf.select %aow6_ok, %aow6_s, %c0 : index + %aow6 = index.assume %aow6_c [range(%aow6_c, 0, 1048575)] : index + %aol6_m = index.mul %vs_tok, %k128 : index + %aol6_a = index.add %aol6_m, %vs_col : index + %aol6_b = index.add %aol6_a, %k1536 : index + %aol6 = index.assume %aol6_b [range(%aol6_b, 0, 2047)] : index + %aov6 = view.load %aost_flat[%aol6] : view<2048xf32> -> f32 + %aog6_m = index.mul %aow6, %attn_stride : index + %aog6_a = index.add %attn_base, %aog6_m : index + %aog6_c = index.add %aog6_a, %vs_gcol : index + %aog6 = index.assume %aog6_c [range(%aog6_c, 0, 1073741823)] : index + scf.if %aow6_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov6, %dst_view, %aog6) : (f32, view<1073741824xf32>, index) + } + %aow7_r = index.add %t0, %k14 : index + %aow7_s = index.add %aow7_r, %vs_tok : index + %aow7_ok = index.cmp ult, %aow7_s, %n_tokens : index + %aow7_c = scf.select %aow7_ok, %aow7_s, %c0 : index + %aow7 = index.assume %aow7_c [range(%aow7_c, 0, 1048575)] : index + %aol7_m = index.mul %vs_tok, %k128 : index + %aol7_a = index.add %aol7_m, %vs_col : index + %aol7_b = index.add %aol7_a, %k1792 : index + %aol7 = index.assume %aol7_b [range(%aol7_b, 0, 2047)] : index + %aov7 = view.load %aost_flat[%aol7] : view<2048xf32> -> f32 + %aog7_m = index.mul %aow7, %attn_stride : index + %aog7_a = index.add %attn_base, %aog7_m : index + %aog7_c = index.add %aog7_a, %vs_gcol : index + %aog7 = index.assume %aog7_c [range(%aog7_c, 0, 1073741823)] : index + scf.if %aow7_ok { + template.apply<@llm.gated_delta_net.publish_f32.scalar>(%aov7, %dst_view, %aog7) : (f32, view<1073741824xf32>, index) + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + // K' was formed from the exact normalized f16 K tile before it + // was overwritten. Only the first 128 threads own rows. + scf.if %kt_live { + %ktrow2 = index.assume %tid [range(%tid, 0, 127)] : index + %ktdst0 = index.mul %ktrow2, %k24 : index + %ktdst = index.assume %ktdst0 [range(%ktdst0, 0, 3064)] : index + vector.store %ktv0, %ka_flat[%ktdst] : vector<8xf16>, view<3072xf16> + %ktdst8_a = index.add %ktdst, %k8 : index + %ktdst8 = index.assume %ktdst8_a [range(%ktdst8_a, 8, 3064)] : index + vector.store %ktv1, %ka_flat[%ktdst8] : vector<8xf16>, view<3072xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %rowh_r = index.rem %row0, %k64 : index + %rowh = index.assume %rowh_r [range(%rowh_r, 0, 60)] : index + %snapshot_next0, %snapshot_next1, %snapshot_next2, %snapshot_next3, %snapshot_next4, %snapshot_next5, %snapshot_next6, %snapshot_next7, %snapshot_next8, %snapshot_next9, %snapshot_next10, %snapshot_next11, %snapshot_next12, %snapshot_next13, %snapshot_next14, %snapshot_next15 = scf.if %reuse_snapshot_fragments -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + scf.if %publish_snapshots { + %snapshot_c0o = index.constant 0 : offset + %snapshot_c0 = index.constant 0 : index + %snapshot_c4 = index.constant 4 : index + %snapshot_c16 = index.constant 16 : index + %snapshot_c32 = index.constant 32 : index + %snapshot_c48 = index.constant 48 : index + %snapshot_c64 = index.constant 64 : index + %snapshot_c68 = index.constant 68 : index + %snapshot_c80 = index.constant 80 : index + %snapshot_c96 = index.constant 96 : index + %snapshot_c112 = index.constant 112 : index + %snapshot_c128 = index.constant 128 : index + %snapshot_c4352 = index.constant 4352 : index + %snapshot_c6528 = index.constant 6528 : index + %snapshot_c11072 = index.constant 11072 : index + %snapshot_ka_o = index.constant 4352 : offset + %snapshot_lay_dcol = encoding.layout.strided [24, 1] : encoding + %snapshot_lay_katok = encoding.layout.strided [1, 24] : encoding + %snapshot_lay_upd = encoding.layout.strided [68, 1] : encoding + %snapshot_lay_state = encoding.layout.strided [128, 1] : encoding + %snapshot_wave_bytes = index.mul %wave, %snapshot_c6528 : index + %snapshot_wave_base = index.add %snapshot_wave_bytes, %snapshot_c11072 : index + %snapshot_wave_off = index.cast %snapshot_wave_base : index to offset + %snapshot_d_base_i = index.add %snapshot_wave_base, %snapshot_c4352 : index + %snapshot_d_base = index.cast %snapshot_d_base_i : index to offset + %snapshot_d_lhs = buffer.view %lds[%snapshot_d_base] : buffer -> view<16x16xf16, %snapshot_lay_dcol> + %snapshot_ka_rhs = buffer.view %lds[%snapshot_ka_o] : buffer -> view<16x128xf16, %snapshot_lay_katok> + %snapshot_upd_flat = buffer.view %lds[%snapshot_wave_off] : buffer -> view<1088xf32> + %snapshot_upd_tile = buffer.view %lds[%snapshot_wave_off] : buffer -> view<16x68xf32, %snapshot_lay_upd> + %snapshot_lane_low = index.cmp ult, %lane, %snapshot_c16 : index + %snapshot_lane_high = index.cmp uge, %lane, %snapshot_c16 : index + scf.if %snapshot_lane_low { + vector.store %sa0, %snapshot_upd_flat[%rowh] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s1_off = index.constant 68 : index + %snapshot_lo_s1_raw = index.add %snapshot_lo_s1_off, %rowh : index + %snapshot_lo_s1_idx = index.assume %snapshot_lo_s1_raw [range(%snapshot_lo_s1_raw, 68, 128)] : index + vector.store %sa1, %snapshot_upd_flat[%snapshot_lo_s1_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s2_off = index.constant 136 : index + %snapshot_lo_s2_raw = index.add %snapshot_lo_s2_off, %rowh : index + %snapshot_lo_s2_idx = index.assume %snapshot_lo_s2_raw [range(%snapshot_lo_s2_raw, 136, 196)] : index + vector.store %sa2, %snapshot_upd_flat[%snapshot_lo_s2_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s3_off = index.constant 204 : index + %snapshot_lo_s3_raw = index.add %snapshot_lo_s3_off, %rowh : index + %snapshot_lo_s3_idx = index.assume %snapshot_lo_s3_raw [range(%snapshot_lo_s3_raw, 204, 264)] : index + vector.store %sa3, %snapshot_upd_flat[%snapshot_lo_s3_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s4_off = index.constant 272 : index + %snapshot_lo_s4_raw = index.add %snapshot_lo_s4_off, %rowh : index + %snapshot_lo_s4_idx = index.assume %snapshot_lo_s4_raw [range(%snapshot_lo_s4_raw, 272, 332)] : index + vector.store %sa4, %snapshot_upd_flat[%snapshot_lo_s4_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s5_off = index.constant 340 : index + %snapshot_lo_s5_raw = index.add %snapshot_lo_s5_off, %rowh : index + %snapshot_lo_s5_idx = index.assume %snapshot_lo_s5_raw [range(%snapshot_lo_s5_raw, 340, 400)] : index + vector.store %sa5, %snapshot_upd_flat[%snapshot_lo_s5_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s6_off = index.constant 408 : index + %snapshot_lo_s6_raw = index.add %snapshot_lo_s6_off, %rowh : index + %snapshot_lo_s6_idx = index.assume %snapshot_lo_s6_raw [range(%snapshot_lo_s6_raw, 408, 468)] : index + vector.store %sa6, %snapshot_upd_flat[%snapshot_lo_s6_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s7_off = index.constant 476 : index + %snapshot_lo_s7_raw = index.add %snapshot_lo_s7_off, %rowh : index + %snapshot_lo_s7_idx = index.assume %snapshot_lo_s7_raw [range(%snapshot_lo_s7_raw, 476, 536)] : index + vector.store %sa7, %snapshot_upd_flat[%snapshot_lo_s7_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s8_off = index.constant 544 : index + %snapshot_lo_s8_raw = index.add %snapshot_lo_s8_off, %rowh : index + %snapshot_lo_s8_idx = index.assume %snapshot_lo_s8_raw [range(%snapshot_lo_s8_raw, 544, 604)] : index + vector.store %sa8, %snapshot_upd_flat[%snapshot_lo_s8_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s9_off = index.constant 612 : index + %snapshot_lo_s9_raw = index.add %snapshot_lo_s9_off, %rowh : index + %snapshot_lo_s9_idx = index.assume %snapshot_lo_s9_raw [range(%snapshot_lo_s9_raw, 612, 672)] : index + vector.store %sa9, %snapshot_upd_flat[%snapshot_lo_s9_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s10_off = index.constant 680 : index + %snapshot_lo_s10_raw = index.add %snapshot_lo_s10_off, %rowh : index + %snapshot_lo_s10_idx = index.assume %snapshot_lo_s10_raw [range(%snapshot_lo_s10_raw, 680, 740)] : index + vector.store %sa10, %snapshot_upd_flat[%snapshot_lo_s10_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s11_off = index.constant 748 : index + %snapshot_lo_s11_raw = index.add %snapshot_lo_s11_off, %rowh : index + %snapshot_lo_s11_idx = index.assume %snapshot_lo_s11_raw [range(%snapshot_lo_s11_raw, 748, 808)] : index + vector.store %sa11, %snapshot_upd_flat[%snapshot_lo_s11_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s12_off = index.constant 816 : index + %snapshot_lo_s12_raw = index.add %snapshot_lo_s12_off, %rowh : index + %snapshot_lo_s12_idx = index.assume %snapshot_lo_s12_raw [range(%snapshot_lo_s12_raw, 816, 876)] : index + vector.store %sa12, %snapshot_upd_flat[%snapshot_lo_s12_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s13_off = index.constant 884 : index + %snapshot_lo_s13_raw = index.add %snapshot_lo_s13_off, %rowh : index + %snapshot_lo_s13_idx = index.assume %snapshot_lo_s13_raw [range(%snapshot_lo_s13_raw, 884, 944)] : index + vector.store %sa13, %snapshot_upd_flat[%snapshot_lo_s13_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s14_off = index.constant 952 : index + %snapshot_lo_s14_raw = index.add %snapshot_lo_s14_off, %rowh : index + %snapshot_lo_s14_idx = index.assume %snapshot_lo_s14_raw [range(%snapshot_lo_s14_raw, 952, 1012)] : index + vector.store %sa14, %snapshot_upd_flat[%snapshot_lo_s14_idx] : vector<4xf32>, view<1088xf32> + %snapshot_lo_s15_off = index.constant 1020 : index + %snapshot_lo_s15_raw = index.add %snapshot_lo_s15_off, %rowh : index + %snapshot_lo_s15_idx = index.assume %snapshot_lo_s15_raw [range(%snapshot_lo_s15_raw, 1020, 1080)] : index + vector.store %sa15, %snapshot_upd_flat[%snapshot_lo_s15_idx] : vector<4xf32>, view<1088xf32> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + %snapshot_lo0_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c0] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + %snapshot_lo1_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c16] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + %snapshot_lo2_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c32] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + %snapshot_lo3_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c48] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %snapshot_lane_high { + vector.store %sa0, %snapshot_upd_flat[%rowh] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s1_off = index.constant 68 : index + %snapshot_hi_s1_raw = index.add %snapshot_hi_s1_off, %rowh : index + %snapshot_hi_s1_idx = index.assume %snapshot_hi_s1_raw [range(%snapshot_hi_s1_raw, 68, 128)] : index + vector.store %sa1, %snapshot_upd_flat[%snapshot_hi_s1_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s2_off = index.constant 136 : index + %snapshot_hi_s2_raw = index.add %snapshot_hi_s2_off, %rowh : index + %snapshot_hi_s2_idx = index.assume %snapshot_hi_s2_raw [range(%snapshot_hi_s2_raw, 136, 196)] : index + vector.store %sa2, %snapshot_upd_flat[%snapshot_hi_s2_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s3_off = index.constant 204 : index + %snapshot_hi_s3_raw = index.add %snapshot_hi_s3_off, %rowh : index + %snapshot_hi_s3_idx = index.assume %snapshot_hi_s3_raw [range(%snapshot_hi_s3_raw, 204, 264)] : index + vector.store %sa3, %snapshot_upd_flat[%snapshot_hi_s3_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s4_off = index.constant 272 : index + %snapshot_hi_s4_raw = index.add %snapshot_hi_s4_off, %rowh : index + %snapshot_hi_s4_idx = index.assume %snapshot_hi_s4_raw [range(%snapshot_hi_s4_raw, 272, 332)] : index + vector.store %sa4, %snapshot_upd_flat[%snapshot_hi_s4_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s5_off = index.constant 340 : index + %snapshot_hi_s5_raw = index.add %snapshot_hi_s5_off, %rowh : index + %snapshot_hi_s5_idx = index.assume %snapshot_hi_s5_raw [range(%snapshot_hi_s5_raw, 340, 400)] : index + vector.store %sa5, %snapshot_upd_flat[%snapshot_hi_s5_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s6_off = index.constant 408 : index + %snapshot_hi_s6_raw = index.add %snapshot_hi_s6_off, %rowh : index + %snapshot_hi_s6_idx = index.assume %snapshot_hi_s6_raw [range(%snapshot_hi_s6_raw, 408, 468)] : index + vector.store %sa6, %snapshot_upd_flat[%snapshot_hi_s6_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s7_off = index.constant 476 : index + %snapshot_hi_s7_raw = index.add %snapshot_hi_s7_off, %rowh : index + %snapshot_hi_s7_idx = index.assume %snapshot_hi_s7_raw [range(%snapshot_hi_s7_raw, 476, 536)] : index + vector.store %sa7, %snapshot_upd_flat[%snapshot_hi_s7_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s8_off = index.constant 544 : index + %snapshot_hi_s8_raw = index.add %snapshot_hi_s8_off, %rowh : index + %snapshot_hi_s8_idx = index.assume %snapshot_hi_s8_raw [range(%snapshot_hi_s8_raw, 544, 604)] : index + vector.store %sa8, %snapshot_upd_flat[%snapshot_hi_s8_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s9_off = index.constant 612 : index + %snapshot_hi_s9_raw = index.add %snapshot_hi_s9_off, %rowh : index + %snapshot_hi_s9_idx = index.assume %snapshot_hi_s9_raw [range(%snapshot_hi_s9_raw, 612, 672)] : index + vector.store %sa9, %snapshot_upd_flat[%snapshot_hi_s9_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s10_off = index.constant 680 : index + %snapshot_hi_s10_raw = index.add %snapshot_hi_s10_off, %rowh : index + %snapshot_hi_s10_idx = index.assume %snapshot_hi_s10_raw [range(%snapshot_hi_s10_raw, 680, 740)] : index + vector.store %sa10, %snapshot_upd_flat[%snapshot_hi_s10_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s11_off = index.constant 748 : index + %snapshot_hi_s11_raw = index.add %snapshot_hi_s11_off, %rowh : index + %snapshot_hi_s11_idx = index.assume %snapshot_hi_s11_raw [range(%snapshot_hi_s11_raw, 748, 808)] : index + vector.store %sa11, %snapshot_upd_flat[%snapshot_hi_s11_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s12_off = index.constant 816 : index + %snapshot_hi_s12_raw = index.add %snapshot_hi_s12_off, %rowh : index + %snapshot_hi_s12_idx = index.assume %snapshot_hi_s12_raw [range(%snapshot_hi_s12_raw, 816, 876)] : index + vector.store %sa12, %snapshot_upd_flat[%snapshot_hi_s12_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s13_off = index.constant 884 : index + %snapshot_hi_s13_raw = index.add %snapshot_hi_s13_off, %rowh : index + %snapshot_hi_s13_idx = index.assume %snapshot_hi_s13_raw [range(%snapshot_hi_s13_raw, 884, 944)] : index + vector.store %sa13, %snapshot_upd_flat[%snapshot_hi_s13_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s14_off = index.constant 952 : index + %snapshot_hi_s14_raw = index.add %snapshot_hi_s14_off, %rowh : index + %snapshot_hi_s14_idx = index.assume %snapshot_hi_s14_raw [range(%snapshot_hi_s14_raw, 952, 1012)] : index + vector.store %sa14, %snapshot_upd_flat[%snapshot_hi_s14_idx] : vector<4xf32>, view<1088xf32> + %snapshot_hi_s15_off = index.constant 1020 : index + %snapshot_hi_s15_raw = index.add %snapshot_hi_s15_off, %rowh : index + %snapshot_hi_s15_idx = index.assume %snapshot_hi_s15_raw [range(%snapshot_hi_s15_raw, 1020, 1080)] : index + vector.store %sa15, %snapshot_upd_flat[%snapshot_hi_s15_idx] : vector<4xf32>, view<1088xf32> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + %snapshot_hi0_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c0] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + %snapshot_hi1_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c16] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + %snapshot_hi2_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c32] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + %snapshot_hi3_state = vector.fragment.load %snapshot_upd_tile[%snapshot_c0, %snapshot_c48] shape [%snapshot_c16, %snapshot_c16] : view<16x68xf32, %snapshot_lay_upd> -> vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + %snapshot_zero = scalar.constant 0.0 : f16 + %snapshot_last = index.sub %n_tokens, %k1 : index + %has_prefix0 = index.cmp ult, %k1, %n_tokens : index + %has_prefix1 = index.cmp ult, %k2, %n_tokens : index + %has_prefix2 = index.cmp ult, %k3, %n_tokens : index + %has_prefix3 = index.cmp ult, %k4, %n_tokens : index + scf.if %has_prefix0 { + %snapshot_slot0 = index.mul %snapshot_stride, %snapshot_last : index + %p0d0 = scalar.fptrunc %dl0 : f32 to f16 + %p0v0 = vector.from_elements %p0d0, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p0v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p0v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p0v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_fragment_prefix.body>(%cp0, %snapshot_slot0, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %snapshot_lo0_state, %snapshot_lo1_state, %snapshot_lo2_state, %snapshot_lo3_state, %snapshot_hi0_state, %snapshot_hi1_state, %snapshot_hi2_state, %snapshot_hi3_state, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %has_prefix1 { + %snapshot_index1 = index.sub %snapshot_last, %k1 : index + %snapshot_slot1 = index.mul %snapshot_stride, %snapshot_index1 : index + %p10 = scalar.divf %cp1, %cp0 : f32 + %p1d0f = scalar.mulf %dl0, %p10 : f32 + %p1d0 = scalar.fptrunc %p1d0f : f32 to f16 + %p1d1 = scalar.fptrunc %dl1 : f32 to f16 + %p1v0 = vector.from_elements %p1d0, %p1d1, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p1v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p1v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p1v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_fragment_prefix.body>(%cp1, %snapshot_slot1, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %snapshot_lo0_state, %snapshot_lo1_state, %snapshot_lo2_state, %snapshot_lo3_state, %snapshot_hi0_state, %snapshot_hi1_state, %snapshot_hi2_state, %snapshot_hi3_state, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %has_prefix2 { + %snapshot_index2 = index.sub %snapshot_last, %k2 : index + %snapshot_slot2 = index.mul %snapshot_stride, %snapshot_index2 : index + %p20 = scalar.divf %cp2, %cp0 : f32 + %p21 = scalar.divf %cp2, %cp1 : f32 + %p2d0f = scalar.mulf %dl0, %p20 : f32 + %p2d1f = scalar.mulf %dl1, %p21 : f32 + %p2d0 = scalar.fptrunc %p2d0f : f32 to f16 + %p2d1 = scalar.fptrunc %p2d1f : f32 to f16 + %p2d2 = scalar.fptrunc %dl2 : f32 to f16 + %p2v0 = vector.from_elements %p2d0, %p2d1, %p2d2, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p2v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p2v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p2v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_fragment_prefix.body>(%cp2, %snapshot_slot2, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %snapshot_lo0_state, %snapshot_lo1_state, %snapshot_lo2_state, %snapshot_lo3_state, %snapshot_hi0_state, %snapshot_hi1_state, %snapshot_hi2_state, %snapshot_hi3_state, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %has_prefix3 { + %snapshot_index3 = index.sub %snapshot_last, %k3 : index + %snapshot_slot3 = index.mul %snapshot_stride, %snapshot_index3 : index + %p30 = scalar.divf %cp3, %cp0 : f32 + %p31 = scalar.divf %cp3, %cp1 : f32 + %p32 = scalar.divf %cp3, %cp2 : f32 + %p3d0f = scalar.mulf %dl0, %p30 : f32 + %p3d1f = scalar.mulf %dl1, %p31 : f32 + %p3d2f = scalar.mulf %dl2, %p32 : f32 + %p3d0 = scalar.fptrunc %p3d0f : f32 to f16 + %p3d1 = scalar.fptrunc %p3d1f : f32 to f16 + %p3d2 = scalar.fptrunc %p3d2f : f32 to f16 + %p3d3 = scalar.fptrunc %dl3 : f32 to f16 + %p3v0 = vector.from_elements %p3d0, %p3d1, %p3d2, %p3d3, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p3v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p3v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p3v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_fragment_prefix.body>(%cp3, %snapshot_slot3, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %snapshot_lo0_state, %snapshot_lo1_state, %snapshot_lo2_state, %snapshot_lo3_state, %snapshot_hi0_state, %snapshot_hi1_state, %snapshot_hi2_state, %snapshot_hi3_state, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + + scf.if %lane_live { + vector.store %ddv0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %ddv1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %snapshot_clast = view.load %gs_flat[%k61] : view<80xf32> -> f32 + template.apply<@llm.gated_delta_net.snapshot_fragment_prefix.body>(%snapshot_clast, %k0, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %snapshot_lo0_state, %snapshot_lo1_state, %snapshot_lo2_state, %snapshot_lo3_state, %snapshot_hi0_state, %snapshot_hi1_state, %snapshot_hi2_state, %snapshot_hi3_state, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) + } + scf.yield %sa0, %sa1, %sa2, %sa3, %sa4, %sa5, %sa6, %sa7, %sa8, %sa9, %sa10, %sa11, %sa12, %sa13, %sa14, %sa15 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } else { + scf.if %publish_snapshots { + %snapshot_zero = scalar.constant 0.0 : f16 + %snapshot_last = index.sub %n_tokens, %k1 : index + %has_prefix0 = index.cmp ult, %k1, %n_tokens : index + %has_prefix1 = index.cmp ult, %k2, %n_tokens : index + %has_prefix2 = index.cmp ult, %k3, %n_tokens : index + %has_prefix3 = index.cmp ult, %k4, %n_tokens : index + scf.if %has_prefix0 { + %snapshot_slot0 = index.mul %snapshot_stride, %snapshot_last : index + %p0d0 = scalar.fptrunc %dl0 : f32 to f16 + %p0v0 = vector.from_elements %p0d0, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p0v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p0v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p0v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_prefix.body>(%cp0, %snapshot_slot0, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %sa0, %sa1, %sa2, %sa3, %sa4, %sa5, %sa6, %sa7, %sa8, %sa9, %sa10, %sa11, %sa12, %sa13, %sa14, %sa15, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %has_prefix1 { + %snapshot_index1 = index.sub %snapshot_last, %k1 : index + %snapshot_slot1 = index.mul %snapshot_stride, %snapshot_index1 : index + %p10 = scalar.divf %cp1, %cp0 : f32 + %p1d0f = scalar.mulf %dl0, %p10 : f32 + %p1d0 = scalar.fptrunc %p1d0f : f32 to f16 + %p1d1 = scalar.fptrunc %dl1 : f32 to f16 + %p1v0 = vector.from_elements %p1d0, %p1d1, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p1v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p1v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p1v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_prefix.body>(%cp1, %snapshot_slot1, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %sa0, %sa1, %sa2, %sa3, %sa4, %sa5, %sa6, %sa7, %sa8, %sa9, %sa10, %sa11, %sa12, %sa13, %sa14, %sa15, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %has_prefix2 { + %snapshot_index2 = index.sub %snapshot_last, %k2 : index + %snapshot_slot2 = index.mul %snapshot_stride, %snapshot_index2 : index + %p20 = scalar.divf %cp2, %cp0 : f32 + %p21 = scalar.divf %cp2, %cp1 : f32 + %p2d0f = scalar.mulf %dl0, %p20 : f32 + %p2d1f = scalar.mulf %dl1, %p21 : f32 + %p2d0 = scalar.fptrunc %p2d0f : f32 to f16 + %p2d1 = scalar.fptrunc %p2d1f : f32 to f16 + %p2d2 = scalar.fptrunc %dl2 : f32 to f16 + %p2v0 = vector.from_elements %p2d0, %p2d1, %p2d2, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p2v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p2v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p2v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_prefix.body>(%cp2, %snapshot_slot2, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %sa0, %sa1, %sa2, %sa3, %sa4, %sa5, %sa6, %sa7, %sa8, %sa9, %sa10, %sa11, %sa12, %sa13, %sa14, %sa15, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.if %has_prefix3 { + %snapshot_index3 = index.sub %snapshot_last, %k3 : index + %snapshot_slot3 = index.mul %snapshot_stride, %snapshot_index3 : index + %p30 = scalar.divf %cp3, %cp0 : f32 + %p31 = scalar.divf %cp3, %cp1 : f32 + %p32 = scalar.divf %cp3, %cp2 : f32 + %p3d0f = scalar.mulf %dl0, %p30 : f32 + %p3d1f = scalar.mulf %dl1, %p31 : f32 + %p3d2f = scalar.mulf %dl2, %p32 : f32 + %p3d0 = scalar.fptrunc %p3d0f : f32 to f16 + %p3d1 = scalar.fptrunc %p3d1f : f32 to f16 + %p3d2 = scalar.fptrunc %p3d2f : f32 to f16 + %p3d3 = scalar.fptrunc %dl3 : f32 to f16 + %p3v0 = vector.from_elements %p3d0, %p3d1, %p3d2, %p3d3, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + %p3v1 = vector.from_elements %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero, %snapshot_zero : vector<8xf16> + scf.if %lane_live { + vector.store %p3v0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %p3v1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@llm.gated_delta_net.snapshot_prefix.body>(%cp3, %snapshot_slot3, %state_offset, %wcol, %row0, %rowh, %wave, %lane, %lds, %snapshot_cache, %sa0, %sa1, %sa2, %sa3, %sa4, %sa5, %sa6, %sa7, %sa8, %sa9, %sa10, %sa11, %sa12, %sa13, %sa14, %sa15, %fz8) : (f32, index, index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<8xf32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + + scf.if %lane_live { + vector.store %ddv0, %d_flat[%dds0] : vector<8xf16>, view<384xf16> + vector.store %ddv1, %d_flat[%dds1] : vector<8xf16>, view<384xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %clast = view.load %gs_flat[%k61] : view<80xf32> -> f32 + %clv = vector.splat %clast : vector<4xf32> + + // Update S <- c_last S + D K' in two 64-row halves and overlay the dead F16 state copy. + // Lane L owns rows 4L..4L+3 in the half selected by its 16-lane group. + %up00_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up00_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up00_r = vector.fragment.load %ka_rhs[%k0, %k0] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up00_m = vector.mma %up00_l, %up00_r, %up00_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + %up01_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up01_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up01_r = vector.fragment.load %ka_rhs[%k0, %k16] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up01_m = vector.mma %up01_l, %up01_r, %up01_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + %up02_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up02_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up02_r = vector.fragment.load %ka_rhs[%k0, %k32] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up02_m = vector.mma %up02_l, %up02_r, %up02_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + %up03_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up03_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up03_r = vector.fragment.load %ka_rhs[%k0, %k48] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up03_m = vector.mma %up03_l, %up03_r, %up03_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %up00_m, %upd_tile[%k0, %k0] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + vector.fragment.store %up01_m, %upd_tile[%k0, %k16] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + vector.fragment.store %up02_m, %upd_tile[%k0, %k32] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + vector.fragment.store %up03_m, %upd_tile[%k0, %k48] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + kernel.barrier scope(subgroup) ordering(acq_rel) + %ur00_m = index.mul %k0, %k68 : index + %ur00_a = index.add %ur00_m, %rowh : index + %ur00_b = index.assume %ur00_a [range(%ur00_a, 0, 1084)] : index + %ur00 = vector.load %upd_flat[%ur00_b] : view<1088xf32> -> vector<4xf32> + %ur01_m = index.mul %k1, %k68 : index + %ur01_a = index.add %ur01_m, %rowh : index + %ur01_b = index.assume %ur01_a [range(%ur01_a, 0, 1084)] : index + %ur01 = vector.load %upd_flat[%ur01_b] : view<1088xf32> -> vector<4xf32> + %ur02_m = index.mul %k2, %k68 : index + %ur02_a = index.add %ur02_m, %rowh : index + %ur02_b = index.assume %ur02_a [range(%ur02_a, 0, 1084)] : index + %ur02 = vector.load %upd_flat[%ur02_b] : view<1088xf32> -> vector<4xf32> + %ur03_m = index.mul %k3, %k68 : index + %ur03_a = index.add %ur03_m, %rowh : index + %ur03_b = index.assume %ur03_a [range(%ur03_a, 0, 1084)] : index + %ur03 = vector.load %upd_flat[%ur03_b] : view<1088xf32> -> vector<4xf32> + %ur04_m = index.mul %k4, %k68 : index + %ur04_a = index.add %ur04_m, %rowh : index + %ur04_b = index.assume %ur04_a [range(%ur04_a, 0, 1084)] : index + %ur04 = vector.load %upd_flat[%ur04_b] : view<1088xf32> -> vector<4xf32> + %ur05_m = index.mul %k5, %k68 : index + %ur05_a = index.add %ur05_m, %rowh : index + %ur05_b = index.assume %ur05_a [range(%ur05_a, 0, 1084)] : index + %ur05 = vector.load %upd_flat[%ur05_b] : view<1088xf32> -> vector<4xf32> + %ur06_m = index.mul %k6, %k68 : index + %ur06_a = index.add %ur06_m, %rowh : index + %ur06_b = index.assume %ur06_a [range(%ur06_a, 0, 1084)] : index + %ur06 = vector.load %upd_flat[%ur06_b] : view<1088xf32> -> vector<4xf32> + %ur07_m = index.mul %k7, %k68 : index + %ur07_a = index.add %ur07_m, %rowh : index + %ur07_b = index.assume %ur07_a [range(%ur07_a, 0, 1084)] : index + %ur07 = vector.load %upd_flat[%ur07_b] : view<1088xf32> -> vector<4xf32> + %ur08_m = index.mul %k8, %k68 : index + %ur08_a = index.add %ur08_m, %rowh : index + %ur08_b = index.assume %ur08_a [range(%ur08_a, 0, 1084)] : index + %ur08 = vector.load %upd_flat[%ur08_b] : view<1088xf32> -> vector<4xf32> + %ur09_m = index.mul %k9, %k68 : index + %ur09_a = index.add %ur09_m, %rowh : index + %ur09_b = index.assume %ur09_a [range(%ur09_a, 0, 1084)] : index + %ur09 = vector.load %upd_flat[%ur09_b] : view<1088xf32> -> vector<4xf32> + %ur010_m = index.mul %k10, %k68 : index + %ur010_a = index.add %ur010_m, %rowh : index + %ur010_b = index.assume %ur010_a [range(%ur010_a, 0, 1084)] : index + %ur010 = vector.load %upd_flat[%ur010_b] : view<1088xf32> -> vector<4xf32> + %ur011_m = index.mul %k11, %k68 : index + %ur011_a = index.add %ur011_m, %rowh : index + %ur011_b = index.assume %ur011_a [range(%ur011_a, 0, 1084)] : index + %ur011 = vector.load %upd_flat[%ur011_b] : view<1088xf32> -> vector<4xf32> + %ur012_m = index.mul %k12, %k68 : index + %ur012_a = index.add %ur012_m, %rowh : index + %ur012_b = index.assume %ur012_a [range(%ur012_a, 0, 1084)] : index + %ur012 = vector.load %upd_flat[%ur012_b] : view<1088xf32> -> vector<4xf32> + %ur013_m = index.mul %k13, %k68 : index + %ur013_a = index.add %ur013_m, %rowh : index + %ur013_b = index.assume %ur013_a [range(%ur013_a, 0, 1084)] : index + %ur013 = vector.load %upd_flat[%ur013_b] : view<1088xf32> -> vector<4xf32> + %ur014_m = index.mul %k14, %k68 : index + %ur014_a = index.add %ur014_m, %rowh : index + %ur014_b = index.assume %ur014_a [range(%ur014_a, 0, 1084)] : index + %ur014 = vector.load %upd_flat[%ur014_b] : view<1088xf32> -> vector<4xf32> + %ur015_m = index.mul %k15, %k68 : index + %ur015_a = index.add %ur015_m, %rowh : index + %ur015_b = index.assume %ur015_a [range(%ur015_a, 0, 1084)] : index + %ur015 = vector.load %upd_flat[%ur015_b] : view<1088xf32> -> vector<4xf32> + %up10_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up10_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up10_r = vector.fragment.load %ka_rhs[%k0, %k64] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up10_m = vector.mma %up10_l, %up10_r, %up10_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + %up11_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up11_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up11_r = vector.fragment.load %ka_rhs[%k0, %k80] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up11_m = vector.mma %up11_l, %up11_r, %up11_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + %up12_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up12_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up12_r = vector.fragment.load %ka_rhs[%k0, %k96] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up12_m = vector.mma %up12_l, %up12_r, %up12_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + %up13_i = vector.fragment %fz8 shape [%k16, %k16] : vector<8xf32> + %up13_l = vector.fragment.load %d_lhs[%k0, %k0] shape [%k16, %k16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %up13_r = vector.fragment.load %ka_rhs[%k0, %k112] shape [%k16, %k16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %up13_m = vector.mma %up13_l, %up13_r, %up13_i : vector<16xf16>, vector<16xf16>, vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %up10_m, %upd_tile[%k0, %k0] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + vector.fragment.store %up11_m, %upd_tile[%k0, %k16] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + vector.fragment.store %up12_m, %upd_tile[%k0, %k32] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + vector.fragment.store %up13_m, %upd_tile[%k0, %k48] shape [%k16, %k16] : vector<8xf32>, view<16x68xf32, %lay_upd> + kernel.barrier scope(subgroup) ordering(acq_rel) + %ur10_m = index.mul %k0, %k68 : index + %ur10_a = index.add %ur10_m, %rowh : index + %ur10_b = index.assume %ur10_a [range(%ur10_a, 0, 1084)] : index + %ur10 = vector.load %upd_flat[%ur10_b] : view<1088xf32> -> vector<4xf32> + %ur11_m = index.mul %k1, %k68 : index + %ur11_a = index.add %ur11_m, %rowh : index + %ur11_b = index.assume %ur11_a [range(%ur11_a, 0, 1084)] : index + %ur11 = vector.load %upd_flat[%ur11_b] : view<1088xf32> -> vector<4xf32> + %ur12_m = index.mul %k2, %k68 : index + %ur12_a = index.add %ur12_m, %rowh : index + %ur12_b = index.assume %ur12_a [range(%ur12_a, 0, 1084)] : index + %ur12 = vector.load %upd_flat[%ur12_b] : view<1088xf32> -> vector<4xf32> + %ur13_m = index.mul %k3, %k68 : index + %ur13_a = index.add %ur13_m, %rowh : index + %ur13_b = index.assume %ur13_a [range(%ur13_a, 0, 1084)] : index + %ur13 = vector.load %upd_flat[%ur13_b] : view<1088xf32> -> vector<4xf32> + %ur14_m = index.mul %k4, %k68 : index + %ur14_a = index.add %ur14_m, %rowh : index + %ur14_b = index.assume %ur14_a [range(%ur14_a, 0, 1084)] : index + %ur14 = vector.load %upd_flat[%ur14_b] : view<1088xf32> -> vector<4xf32> + %ur15_m = index.mul %k5, %k68 : index + %ur15_a = index.add %ur15_m, %rowh : index + %ur15_b = index.assume %ur15_a [range(%ur15_a, 0, 1084)] : index + %ur15 = vector.load %upd_flat[%ur15_b] : view<1088xf32> -> vector<4xf32> + %ur16_m = index.mul %k6, %k68 : index + %ur16_a = index.add %ur16_m, %rowh : index + %ur16_b = index.assume %ur16_a [range(%ur16_a, 0, 1084)] : index + %ur16 = vector.load %upd_flat[%ur16_b] : view<1088xf32> -> vector<4xf32> + %ur17_m = index.mul %k7, %k68 : index + %ur17_a = index.add %ur17_m, %rowh : index + %ur17_b = index.assume %ur17_a [range(%ur17_a, 0, 1084)] : index + %ur17 = vector.load %upd_flat[%ur17_b] : view<1088xf32> -> vector<4xf32> + %ur18_m = index.mul %k8, %k68 : index + %ur18_a = index.add %ur18_m, %rowh : index + %ur18_b = index.assume %ur18_a [range(%ur18_a, 0, 1084)] : index + %ur18 = vector.load %upd_flat[%ur18_b] : view<1088xf32> -> vector<4xf32> + %ur19_m = index.mul %k9, %k68 : index + %ur19_a = index.add %ur19_m, %rowh : index + %ur19_b = index.assume %ur19_a [range(%ur19_a, 0, 1084)] : index + %ur19 = vector.load %upd_flat[%ur19_b] : view<1088xf32> -> vector<4xf32> + %ur110_m = index.mul %k10, %k68 : index + %ur110_a = index.add %ur110_m, %rowh : index + %ur110_b = index.assume %ur110_a [range(%ur110_a, 0, 1084)] : index + %ur110 = vector.load %upd_flat[%ur110_b] : view<1088xf32> -> vector<4xf32> + %ur111_m = index.mul %k11, %k68 : index + %ur111_a = index.add %ur111_m, %rowh : index + %ur111_b = index.assume %ur111_a [range(%ur111_a, 0, 1084)] : index + %ur111 = vector.load %upd_flat[%ur111_b] : view<1088xf32> -> vector<4xf32> + %ur112_m = index.mul %k12, %k68 : index + %ur112_a = index.add %ur112_m, %rowh : index + %ur112_b = index.assume %ur112_a [range(%ur112_a, 0, 1084)] : index + %ur112 = vector.load %upd_flat[%ur112_b] : view<1088xf32> -> vector<4xf32> + %ur113_m = index.mul %k13, %k68 : index + %ur113_a = index.add %ur113_m, %rowh : index + %ur113_b = index.assume %ur113_a [range(%ur113_a, 0, 1084)] : index + %ur113 = vector.load %upd_flat[%ur113_b] : view<1088xf32> -> vector<4xf32> + %ur114_m = index.mul %k14, %k68 : index + %ur114_a = index.add %ur114_m, %rowh : index + %ur114_b = index.assume %ur114_a [range(%ur114_a, 0, 1084)] : index + %ur114 = vector.load %upd_flat[%ur114_b] : view<1088xf32> -> vector<4xf32> + %ur115_m = index.mul %k15, %k68 : index + %ur115_a = index.add %ur115_m, %rowh : index + %ur115_b = index.assume %ur115_a [range(%ur115_a, 0, 1084)] : index + %ur115 = vector.load %upd_flat[%ur115_b] : view<1088xf32> -> vector<4xf32> + + %usel0 = scf.select %lane_lo, %ur00, %ur10 : vector<4xf32> + %sdec0 = vector.mulf %sa0, %clv : vector<4xf32> + %snew0 = vector.addf %sdec0, %usel0 : vector<4xf32> + %usel1 = scf.select %lane_lo, %ur01, %ur11 : vector<4xf32> + %sdec1 = vector.mulf %sa1, %clv : vector<4xf32> + %snew1 = vector.addf %sdec1, %usel1 : vector<4xf32> + %usel2 = scf.select %lane_lo, %ur02, %ur12 : vector<4xf32> + %sdec2 = vector.mulf %sa2, %clv : vector<4xf32> + %snew2 = vector.addf %sdec2, %usel2 : vector<4xf32> + %usel3 = scf.select %lane_lo, %ur03, %ur13 : vector<4xf32> + %sdec3 = vector.mulf %sa3, %clv : vector<4xf32> + %snew3 = vector.addf %sdec3, %usel3 : vector<4xf32> + %usel4 = scf.select %lane_lo, %ur04, %ur14 : vector<4xf32> + %sdec4 = vector.mulf %sa4, %clv : vector<4xf32> + %snew4 = vector.addf %sdec4, %usel4 : vector<4xf32> + %usel5 = scf.select %lane_lo, %ur05, %ur15 : vector<4xf32> + %sdec5 = vector.mulf %sa5, %clv : vector<4xf32> + %snew5 = vector.addf %sdec5, %usel5 : vector<4xf32> + %usel6 = scf.select %lane_lo, %ur06, %ur16 : vector<4xf32> + %sdec6 = vector.mulf %sa6, %clv : vector<4xf32> + %snew6 = vector.addf %sdec6, %usel6 : vector<4xf32> + %usel7 = scf.select %lane_lo, %ur07, %ur17 : vector<4xf32> + %sdec7 = vector.mulf %sa7, %clv : vector<4xf32> + %snew7 = vector.addf %sdec7, %usel7 : vector<4xf32> + %usel8 = scf.select %lane_lo, %ur08, %ur18 : vector<4xf32> + %sdec8 = vector.mulf %sa8, %clv : vector<4xf32> + %snew8 = vector.addf %sdec8, %usel8 : vector<4xf32> + %usel9 = scf.select %lane_lo, %ur09, %ur19 : vector<4xf32> + %sdec9 = vector.mulf %sa9, %clv : vector<4xf32> + %snew9 = vector.addf %sdec9, %usel9 : vector<4xf32> + %usel10 = scf.select %lane_lo, %ur010, %ur110 : vector<4xf32> + %sdec10 = vector.mulf %sa10, %clv : vector<4xf32> + %snew10 = vector.addf %sdec10, %usel10 : vector<4xf32> + %usel11 = scf.select %lane_lo, %ur011, %ur111 : vector<4xf32> + %sdec11 = vector.mulf %sa11, %clv : vector<4xf32> + %snew11 = vector.addf %sdec11, %usel11 : vector<4xf32> + %usel12 = scf.select %lane_lo, %ur012, %ur112 : vector<4xf32> + %sdec12 = vector.mulf %sa12, %clv : vector<4xf32> + %snew12 = vector.addf %sdec12, %usel12 : vector<4xf32> + %usel13 = scf.select %lane_lo, %ur013, %ur113 : vector<4xf32> + %sdec13 = vector.mulf %sa13, %clv : vector<4xf32> + %snew13 = vector.addf %sdec13, %usel13 : vector<4xf32> + %usel14 = scf.select %lane_lo, %ur014, %ur114 : vector<4xf32> + %sdec14 = vector.mulf %sa14, %clv : vector<4xf32> + %snew14 = vector.addf %sdec14, %usel14 : vector<4xf32> + %usel15 = scf.select %lane_lo, %ur015, %ur115 : vector<4xf32> + %sdec15 = vector.mulf %sa15, %clv : vector<4xf32> + %snew15 = vector.addf %sdec15, %usel15 : vector<4xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %snew0, %snew1, %snew2, %snew3, %snew4, %snew5, %snew6, %snew7, %snew8, %snew9, %snew10, %snew11, %snew12, %snew13, %snew14, %snew15 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %snapshot_next0, %snapshot_next1, %snapshot_next2, %snapshot_next3, %snapshot_next4, %snapshot_next5, %snapshot_next6, %snapshot_next7, %snapshot_next8, %snapshot_next9, %snapshot_next10, %snapshot_next11, %snapshot_next12, %snapshot_next13, %snapshot_next14, %snapshot_next15 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + + scf.if %reuse_snapshot_fragments { + } else { + scf.if %publish_snapshots { + template.apply<@llm.gated_delta_net.snapshot_publish.final_head128>(%snapshot_cache_view, %icb0, %icb1, %icb2, %icb3, %icb4, %icb5, %icb6, %icb7, %icb8, %icb9, %icb10, %icb11, %icb12, %icb13, %icb14, %icb15, %sf0, %sf1, %sf2, %sf3, %sf4, %sf5, %sf6, %sf7, %sf8, %sf9, %sf10, %sf11, %sf12, %sf13, %sf14, %sf15) : (view<1073741824xf32>, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + } else { + scf.if %state_inplace { + // All state shards have already been consumed, and each workgroup owns a + // disjoint head/column range, so publication can safely reuse the cache row. + template.apply<@llm.gated_delta_net.inplace_state.publish_head128>(%si_view, %icb0, %icb1, %icb2, %icb3, %icb4, %icb5, %icb6, %icb7, %icb8, %icb9, %icb10, %icb11, %icb12, %icb13, %icb14, %icb15, %sf0, %sf1, %sf2, %sf3, %sf4, %sf5, %sf6, %sf7, %sf8, %sf9, %sf10, %sf11, %sf12, %sf13, %sf14, %sf15) : (view<1073741824xf32>, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + } else { + // The ordinary graph contract appends updated state after attention. + template.apply<@llm.gated_delta_net.publish_f32.state_head128>(%wcol, %s_v, %state_offset, %attn_elems, %row0, %k128, %dst_view, %sf0, %sf1, %sf2, %sf3, %sf4, %sf5, %sf6, %sf7, %sf8, %sf9, %sf10, %sf11, %sf12, %sf13, %sf14, %sf15) : (index, index, index, index, index, index, view<1073741824xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + } + } + } + scf.if %publish_q8 { + %t0 = index.constant 0 : index + template.apply<@llm.gated_delta_net.rmsnorm_gate.publish_head128>(%base, %n_tokens, %n_heads, %n_seqs, %lane, %wave, %t0, %seq, %c0, %h_idx, %k4, %k8, %k16, %k128, %rms_epsilon, %rms_gate_op, %publish_q8, %lds, %rms_weight, %raw_gate, %norm_output, %half_output) : (offset, index, index, index, index, index, index, index, index, index, index, index, index, index, f32, index, i1, buffer, buffer, buffer, buffer, buffer) + + } + + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/selected_snapshot_load.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/selected_snapshot_load.loom new file mode 100644 index 000000000000..09255254cc36 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/selected_snapshot_load.loom @@ -0,0 +1,27 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.selected_snapshot_load.state_offset(%select_state: i1, %state_row_count: index, %state_ids: buffer, %base: offset, %sequence_count: index, %sequence: index, %head_count: index, %head: index, %state_size: index, %default_offset: index) -> (index, i1) + +template.def<@llm.gated_delta_net.selected_snapshot_load.state_offset> device @llm_gated_delta_net_selected_snapshot_load_state_offset(%select_state: i1, %state_row_count: index, %state_ids: buffer, %base: offset, %sequence_count: index, %sequence: index, %head_count: index, %head: index, %state_size: index, %default_offset: index) -> (index, i1) { + %selected_state_offset, %selected_valid = scf.if %select_state -> (index, i1) { + %selected_zero_i32 = scalar.constant 0 : i32 + %selected_limit_i32 = index.cast %state_row_count : index to i32 + %selected_ids_view = buffer.view %state_ids[%base] : buffer -> view<[%sequence_count]xi32> + %selected_raw = view.load %selected_ids_view[%sequence] : view<[%sequence_count]xi32> -> i32 + %selected_nonnegative = scalar.cmpi sge, %selected_raw, %selected_zero_i32 : i32 + %selected_in_range = scalar.cmpi slt, %selected_raw, %selected_limit_i32 : i32 + %selected_ok = scalar.andi %selected_nonnegative, %selected_in_range : i1 + %selected_safe_i32 = scf.select %selected_ok, %selected_raw, %selected_zero_i32 : i32 + %selected_id0 = index.cast %selected_safe_i32 : i32 to index + %selected_id = index.assume %selected_id0 [range(%selected_id0, 0, 262207), lt(%selected_id0, %state_row_count)] : index + %selected_head0 = index.mul %selected_id, %head_count : index + %selected_head = index.add %selected_head0, %head : index + %selected_offset = index.mul %selected_head, %state_size : index + scf.yield %selected_offset, %selected_ok : index, i1 + } else { + %selected_always_valid = scalar.constant true : i1 + scf.yield %default_offset, %selected_always_valid : index, i1 + } + template.return %selected_state_offset, %selected_valid : index, i1 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/snapshot_load.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/snapshot_load.loom new file mode 100644 index 000000000000..24f281ab1877 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/snapshot_load.loom @@ -0,0 +1,12 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.snapshot_load.vector4(%state: view<1073741824xf32>, %index: index, %state_offset: index, %selected_state_offset: index, %selected_valid: i1, %zero: vector<4xf32>) -> (vector<4xf32>) + +template.def<@llm.gated_delta_net.snapshot_load.vector4> device @llm_gated_delta_net_snapshot_load_vector4(%state: view<1073741824xf32>, %index: index, %state_offset: index, %selected_state_offset: index, %selected_valid: i1, %zero: vector<4xf32>) -> (vector<4xf32>) { + %selected_index_local = index.sub %index, %state_offset : index + %selected_index = index.add %selected_index_local, %selected_state_offset : index + %selected = vector.load %state[%selected_index] : view<1073741824xf32> -> vector<4xf32> + %value = scf.select %selected_valid, %selected, %zero : vector<4xf32> + template.return %value : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/snapshot_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/snapshot_publish.loom new file mode 100644 index 000000000000..92ad7445785b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/gated_delta_net/snapshot_publish.loom @@ -0,0 +1,369 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.snapshot_prefix.body(%state_factor: f32, %cache_base: index, %state_offset: index, %wcol: index, %row0: index, %rowh: index, %wave: index, %lane: index, %lds: buffer, %snapshot_cache: buffer, %sa0: vector<4xf32>, %sa1: vector<4xf32>, %sa2: vector<4xf32>, %sa3: vector<4xf32>, %sa4: vector<4xf32>, %sa5: vector<4xf32>, %sa6: vector<4xf32>, %sa7: vector<4xf32>, %sa8: vector<4xf32>, %sa9: vector<4xf32>, %sa10: vector<4xf32>, %sa11: vector<4xf32>, %sa12: vector<4xf32>, %sa13: vector<4xf32>, %sa14: vector<4xf32>, %sa15: vector<4xf32>, %fz8: vector<8xf32>) + +template.decl @llm.gated_delta_net.snapshot_fragment_prefix.body(%state_factor: f32, %cache_base: index, %state_offset: index, %wcol: index, %row0: index, %rowh: index, %wave: index, %lane: index, %lds: buffer, %snapshot_cache: buffer, %lo0_state: vector<8xf32>, %lo1_state: vector<8xf32>, %lo2_state: vector<8xf32>, %lo3_state: vector<8xf32>, %hi0_state: vector<8xf32>, %hi1_state: vector<8xf32>, %hi2_state: vector<8xf32>, %hi3_state: vector<8xf32>, %fz8: vector<8xf32>) + +template.decl @llm.gated_delta_net.snapshot_publish.final_head128(%cache: view<1073741824xf32>, %index0: index, %index1: index, %index2: index, %index3: index, %index4: index, %index5: index, %index6: index, %index7: index, %index8: index, %index9: index, %index10: index, %index11: index, %index12: index, %index13: index, %index14: index, %index15: index, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) + +template.def<@llm.gated_delta_net.snapshot_publish.final_head128> device @llm_gated_delta_net_snapshot_publish_final_head128(%cache: view<1073741824xf32>, %index0: index, %index1: index, %index2: index, %index3: index, %index4: index, %index5: index, %index6: index, %index7: index, %index8: index, %index9: index, %index10: index, %index11: index, %index12: index, %index13: index, %index14: index, %index15: index, %value0: vector<4xf32>, %value1: vector<4xf32>, %value2: vector<4xf32>, %value3: vector<4xf32>, %value4: vector<4xf32>, %value5: vector<4xf32>, %value6: vector<4xf32>, %value7: vector<4xf32>, %value8: vector<4xf32>, %value9: vector<4xf32>, %value10: vector<4xf32>, %value11: vector<4xf32>, %value12: vector<4xf32>, %value13: vector<4xf32>, %value14: vector<4xf32>, %value15: vector<4xf32>) { + vector.store %value0, %cache[%index0] : vector<4xf32>, view<1073741824xf32> + vector.store %value1, %cache[%index1] : vector<4xf32>, view<1073741824xf32> + vector.store %value2, %cache[%index2] : vector<4xf32>, view<1073741824xf32> + vector.store %value3, %cache[%index3] : vector<4xf32>, view<1073741824xf32> + vector.store %value4, %cache[%index4] : vector<4xf32>, view<1073741824xf32> + vector.store %value5, %cache[%index5] : vector<4xf32>, view<1073741824xf32> + vector.store %value6, %cache[%index6] : vector<4xf32>, view<1073741824xf32> + vector.store %value7, %cache[%index7] : vector<4xf32>, view<1073741824xf32> + vector.store %value8, %cache[%index8] : vector<4xf32>, view<1073741824xf32> + vector.store %value9, %cache[%index9] : vector<4xf32>, view<1073741824xf32> + vector.store %value10, %cache[%index10] : vector<4xf32>, view<1073741824xf32> + vector.store %value11, %cache[%index11] : vector<4xf32>, view<1073741824xf32> + vector.store %value12, %cache[%index12] : vector<4xf32>, view<1073741824xf32> + vector.store %value13, %cache[%index13] : vector<4xf32>, view<1073741824xf32> + vector.store %value14, %cache[%index14] : vector<4xf32>, view<1073741824xf32> + vector.store %value15, %cache[%index15] : vector<4xf32>, view<1073741824xf32> + template.return +} + +// Prefix-state publisher shared by the four-row rollback specialization. +// Stage each resident state half in logical result layout, then combine the +// exact WMMA result in registers and publish it directly to the rollback cache. +template.def<@llm.gated_delta_net.snapshot_prefix.body> device @llm_gated_delta_net_snapshot_prefix_body(%state_factor: f32, %cache_base: index, %state_offset: index, %wcol: index, %row0: index, %rowh: index, %wave: index, %lane: index, %lds: buffer, %snapshot_cache: buffer, %sa0: vector<4xf32>, %sa1: vector<4xf32>, %sa2: vector<4xf32>, %sa3: vector<4xf32>, %sa4: vector<4xf32>, %sa5: vector<4xf32>, %sa6: vector<4xf32>, %sa7: vector<4xf32>, %sa8: vector<4xf32>, %sa9: vector<4xf32>, %sa10: vector<4xf32>, %sa11: vector<4xf32>, %sa12: vector<4xf32>, %sa13: vector<4xf32>, %sa14: vector<4xf32>, %sa15: vector<4xf32>, %fz8: vector<8xf32>) { + %c0o = index.constant 0 : offset + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c48 = index.constant 48 : index + %c64 = index.constant 64 : index + %c68 = index.constant 68 : index + %c80 = index.constant 80 : index + %c96 = index.constant 96 : index + %c112 = index.constant 112 : index + %c128 = index.constant 128 : index + %c4352 = index.constant 4352 : index + %c6528 = index.constant 6528 : index + %c11072 = index.constant 11072 : index + %ka_o = index.constant 4352 : offset + %lay_dcol = encoding.layout.strided [24, 1] : encoding + %lay_katok = encoding.layout.strided [1, 24] : encoding + %lay_upd = encoding.layout.strided [68, 1] : encoding + %lay_state = encoding.layout.strided [128, 1] : encoding + %wave_bytes = index.mul %wave, %c6528 : index + %wave_base = index.add %wave_bytes, %c11072 : index + %wave_off = index.cast %wave_base : index to offset + %d_base_i = index.add %wave_base, %c4352 : index + %d_base = index.cast %d_base_i : index to offset + %d_lhs = buffer.view %lds[%d_base] : buffer -> view<16x16xf16, %lay_dcol> + %ka_rhs = buffer.view %lds[%ka_o] : buffer -> view<16x128xf16, %lay_katok> + %upd_flat = buffer.view %lds[%wave_off] : buffer -> view<1088xf32> + %upd_tile = buffer.view %lds[%wave_off] : buffer -> view<16x68xf32, %lay_upd> + %cache_g = buffer.assume.memory_space %snapshot_cache : buffer + %cache_state_base = index.add %cache_base, %state_offset : index + %wave_state = index.mul %wcol, %c128 : index + %cache_wave_base = index.add %cache_state_base, %wave_state : index + %cache_byte_i = index.mul %cache_wave_base, %c4 : index + %cache_byte = index.cast %cache_byte_i : index to offset + %cache_view = buffer.view %cache_g[%cache_byte] : buffer -> view<16x128xf32, %lay_state> + %state_factor_v = vector.splat %state_factor : vector<8xf32> + %lane_low = index.cmp ult, %lane, %c16 : index + %lane_high = index.cmp uge, %lane, %c16 : index + + scf.if %lane_low { + vector.store %sa0, %upd_flat[%rowh] : vector<4xf32>, view<1088xf32> + %lo_s1_off = index.constant 68 : index + %lo_s1_raw = index.add %lo_s1_off, %rowh : index + %lo_s1_idx = index.assume %lo_s1_raw [range(%lo_s1_raw, 68, 128)] : index + vector.store %sa1, %upd_flat[%lo_s1_idx] : vector<4xf32>, view<1088xf32> + %lo_s2_off = index.constant 136 : index + %lo_s2_raw = index.add %lo_s2_off, %rowh : index + %lo_s2_idx = index.assume %lo_s2_raw [range(%lo_s2_raw, 136, 196)] : index + vector.store %sa2, %upd_flat[%lo_s2_idx] : vector<4xf32>, view<1088xf32> + %lo_s3_off = index.constant 204 : index + %lo_s3_raw = index.add %lo_s3_off, %rowh : index + %lo_s3_idx = index.assume %lo_s3_raw [range(%lo_s3_raw, 204, 264)] : index + vector.store %sa3, %upd_flat[%lo_s3_idx] : vector<4xf32>, view<1088xf32> + %lo_s4_off = index.constant 272 : index + %lo_s4_raw = index.add %lo_s4_off, %rowh : index + %lo_s4_idx = index.assume %lo_s4_raw [range(%lo_s4_raw, 272, 332)] : index + vector.store %sa4, %upd_flat[%lo_s4_idx] : vector<4xf32>, view<1088xf32> + %lo_s5_off = index.constant 340 : index + %lo_s5_raw = index.add %lo_s5_off, %rowh : index + %lo_s5_idx = index.assume %lo_s5_raw [range(%lo_s5_raw, 340, 400)] : index + vector.store %sa5, %upd_flat[%lo_s5_idx] : vector<4xf32>, view<1088xf32> + %lo_s6_off = index.constant 408 : index + %lo_s6_raw = index.add %lo_s6_off, %rowh : index + %lo_s6_idx = index.assume %lo_s6_raw [range(%lo_s6_raw, 408, 468)] : index + vector.store %sa6, %upd_flat[%lo_s6_idx] : vector<4xf32>, view<1088xf32> + %lo_s7_off = index.constant 476 : index + %lo_s7_raw = index.add %lo_s7_off, %rowh : index + %lo_s7_idx = index.assume %lo_s7_raw [range(%lo_s7_raw, 476, 536)] : index + vector.store %sa7, %upd_flat[%lo_s7_idx] : vector<4xf32>, view<1088xf32> + %lo_s8_off = index.constant 544 : index + %lo_s8_raw = index.add %lo_s8_off, %rowh : index + %lo_s8_idx = index.assume %lo_s8_raw [range(%lo_s8_raw, 544, 604)] : index + vector.store %sa8, %upd_flat[%lo_s8_idx] : vector<4xf32>, view<1088xf32> + %lo_s9_off = index.constant 612 : index + %lo_s9_raw = index.add %lo_s9_off, %rowh : index + %lo_s9_idx = index.assume %lo_s9_raw [range(%lo_s9_raw, 612, 672)] : index + vector.store %sa9, %upd_flat[%lo_s9_idx] : vector<4xf32>, view<1088xf32> + %lo_s10_off = index.constant 680 : index + %lo_s10_raw = index.add %lo_s10_off, %rowh : index + %lo_s10_idx = index.assume %lo_s10_raw [range(%lo_s10_raw, 680, 740)] : index + vector.store %sa10, %upd_flat[%lo_s10_idx] : vector<4xf32>, view<1088xf32> + %lo_s11_off = index.constant 748 : index + %lo_s11_raw = index.add %lo_s11_off, %rowh : index + %lo_s11_idx = index.assume %lo_s11_raw [range(%lo_s11_raw, 748, 808)] : index + vector.store %sa11, %upd_flat[%lo_s11_idx] : vector<4xf32>, view<1088xf32> + %lo_s12_off = index.constant 816 : index + %lo_s12_raw = index.add %lo_s12_off, %rowh : index + %lo_s12_idx = index.assume %lo_s12_raw [range(%lo_s12_raw, 816, 876)] : index + vector.store %sa12, %upd_flat[%lo_s12_idx] : vector<4xf32>, view<1088xf32> + %lo_s13_off = index.constant 884 : index + %lo_s13_raw = index.add %lo_s13_off, %rowh : index + %lo_s13_idx = index.assume %lo_s13_raw [range(%lo_s13_raw, 884, 944)] : index + vector.store %sa13, %upd_flat[%lo_s13_idx] : vector<4xf32>, view<1088xf32> + %lo_s14_off = index.constant 952 : index + %lo_s14_raw = index.add %lo_s14_off, %rowh : index + %lo_s14_idx = index.assume %lo_s14_raw [range(%lo_s14_raw, 952, 1012)] : index + vector.store %sa14, %upd_flat[%lo_s14_idx] : vector<4xf32>, view<1088xf32> + %lo_s15_off = index.constant 1020 : index + %lo_s15_raw = index.add %lo_s15_off, %rowh : index + %lo_s15_idx = index.assume %lo_s15_raw [range(%lo_s15_raw, 1020, 1080)] : index + vector.store %sa15, %upd_flat[%lo_s15_idx] : vector<4xf32>, view<1088xf32> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + %lo0_state = vector.fragment.load %upd_tile[%c0, %c0] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %lo0_scaled = vector.mulf %lo0_state, %state_factor_v : vector<8xf32> + %lo0_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo0_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo0_rhs = vector.fragment.load %ka_rhs[%c0, %c0] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo0_mma = vector.mma %lo0_lhs, %lo0_rhs, %lo0_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo0_result = vector.addf %lo0_scaled, %lo0_mma : vector<8xf32> + vector.fragment.store %lo0_result, %cache_view[%c0, %c0] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %lo1_state = vector.fragment.load %upd_tile[%c0, %c16] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %lo1_scaled = vector.mulf %lo1_state, %state_factor_v : vector<8xf32> + %lo1_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo1_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo1_rhs = vector.fragment.load %ka_rhs[%c0, %c16] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo1_mma = vector.mma %lo1_lhs, %lo1_rhs, %lo1_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo1_result = vector.addf %lo1_scaled, %lo1_mma : vector<8xf32> + vector.fragment.store %lo1_result, %cache_view[%c0, %c16] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %lo2_state = vector.fragment.load %upd_tile[%c0, %c32] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %lo2_scaled = vector.mulf %lo2_state, %state_factor_v : vector<8xf32> + %lo2_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo2_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo2_rhs = vector.fragment.load %ka_rhs[%c0, %c32] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo2_mma = vector.mma %lo2_lhs, %lo2_rhs, %lo2_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo2_result = vector.addf %lo2_scaled, %lo2_mma : vector<8xf32> + vector.fragment.store %lo2_result, %cache_view[%c0, %c32] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %lo3_state = vector.fragment.load %upd_tile[%c0, %c48] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %lo3_scaled = vector.mulf %lo3_state, %state_factor_v : vector<8xf32> + %lo3_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo3_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo3_rhs = vector.fragment.load %ka_rhs[%c0, %c48] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo3_mma = vector.mma %lo3_lhs, %lo3_rhs, %lo3_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo3_result = vector.addf %lo3_scaled, %lo3_mma : vector<8xf32> + vector.fragment.store %lo3_result, %cache_view[%c0, %c48] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %lane_high { + vector.store %sa0, %upd_flat[%rowh] : vector<4xf32>, view<1088xf32> + %hi_s1_off = index.constant 68 : index + %hi_s1_raw = index.add %hi_s1_off, %rowh : index + %hi_s1_idx = index.assume %hi_s1_raw [range(%hi_s1_raw, 68, 128)] : index + vector.store %sa1, %upd_flat[%hi_s1_idx] : vector<4xf32>, view<1088xf32> + %hi_s2_off = index.constant 136 : index + %hi_s2_raw = index.add %hi_s2_off, %rowh : index + %hi_s2_idx = index.assume %hi_s2_raw [range(%hi_s2_raw, 136, 196)] : index + vector.store %sa2, %upd_flat[%hi_s2_idx] : vector<4xf32>, view<1088xf32> + %hi_s3_off = index.constant 204 : index + %hi_s3_raw = index.add %hi_s3_off, %rowh : index + %hi_s3_idx = index.assume %hi_s3_raw [range(%hi_s3_raw, 204, 264)] : index + vector.store %sa3, %upd_flat[%hi_s3_idx] : vector<4xf32>, view<1088xf32> + %hi_s4_off = index.constant 272 : index + %hi_s4_raw = index.add %hi_s4_off, %rowh : index + %hi_s4_idx = index.assume %hi_s4_raw [range(%hi_s4_raw, 272, 332)] : index + vector.store %sa4, %upd_flat[%hi_s4_idx] : vector<4xf32>, view<1088xf32> + %hi_s5_off = index.constant 340 : index + %hi_s5_raw = index.add %hi_s5_off, %rowh : index + %hi_s5_idx = index.assume %hi_s5_raw [range(%hi_s5_raw, 340, 400)] : index + vector.store %sa5, %upd_flat[%hi_s5_idx] : vector<4xf32>, view<1088xf32> + %hi_s6_off = index.constant 408 : index + %hi_s6_raw = index.add %hi_s6_off, %rowh : index + %hi_s6_idx = index.assume %hi_s6_raw [range(%hi_s6_raw, 408, 468)] : index + vector.store %sa6, %upd_flat[%hi_s6_idx] : vector<4xf32>, view<1088xf32> + %hi_s7_off = index.constant 476 : index + %hi_s7_raw = index.add %hi_s7_off, %rowh : index + %hi_s7_idx = index.assume %hi_s7_raw [range(%hi_s7_raw, 476, 536)] : index + vector.store %sa7, %upd_flat[%hi_s7_idx] : vector<4xf32>, view<1088xf32> + %hi_s8_off = index.constant 544 : index + %hi_s8_raw = index.add %hi_s8_off, %rowh : index + %hi_s8_idx = index.assume %hi_s8_raw [range(%hi_s8_raw, 544, 604)] : index + vector.store %sa8, %upd_flat[%hi_s8_idx] : vector<4xf32>, view<1088xf32> + %hi_s9_off = index.constant 612 : index + %hi_s9_raw = index.add %hi_s9_off, %rowh : index + %hi_s9_idx = index.assume %hi_s9_raw [range(%hi_s9_raw, 612, 672)] : index + vector.store %sa9, %upd_flat[%hi_s9_idx] : vector<4xf32>, view<1088xf32> + %hi_s10_off = index.constant 680 : index + %hi_s10_raw = index.add %hi_s10_off, %rowh : index + %hi_s10_idx = index.assume %hi_s10_raw [range(%hi_s10_raw, 680, 740)] : index + vector.store %sa10, %upd_flat[%hi_s10_idx] : vector<4xf32>, view<1088xf32> + %hi_s11_off = index.constant 748 : index + %hi_s11_raw = index.add %hi_s11_off, %rowh : index + %hi_s11_idx = index.assume %hi_s11_raw [range(%hi_s11_raw, 748, 808)] : index + vector.store %sa11, %upd_flat[%hi_s11_idx] : vector<4xf32>, view<1088xf32> + %hi_s12_off = index.constant 816 : index + %hi_s12_raw = index.add %hi_s12_off, %rowh : index + %hi_s12_idx = index.assume %hi_s12_raw [range(%hi_s12_raw, 816, 876)] : index + vector.store %sa12, %upd_flat[%hi_s12_idx] : vector<4xf32>, view<1088xf32> + %hi_s13_off = index.constant 884 : index + %hi_s13_raw = index.add %hi_s13_off, %rowh : index + %hi_s13_idx = index.assume %hi_s13_raw [range(%hi_s13_raw, 884, 944)] : index + vector.store %sa13, %upd_flat[%hi_s13_idx] : vector<4xf32>, view<1088xf32> + %hi_s14_off = index.constant 952 : index + %hi_s14_raw = index.add %hi_s14_off, %rowh : index + %hi_s14_idx = index.assume %hi_s14_raw [range(%hi_s14_raw, 952, 1012)] : index + vector.store %sa14, %upd_flat[%hi_s14_idx] : vector<4xf32>, view<1088xf32> + %hi_s15_off = index.constant 1020 : index + %hi_s15_raw = index.add %hi_s15_off, %rowh : index + %hi_s15_idx = index.assume %hi_s15_raw [range(%hi_s15_raw, 1020, 1080)] : index + vector.store %sa15, %upd_flat[%hi_s15_idx] : vector<4xf32>, view<1088xf32> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + %hi0_state = vector.fragment.load %upd_tile[%c0, %c0] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %hi0_scaled = vector.mulf %hi0_state, %state_factor_v : vector<8xf32> + %hi0_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi0_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi0_rhs = vector.fragment.load %ka_rhs[%c0, %c64] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi0_mma = vector.mma %hi0_lhs, %hi0_rhs, %hi0_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi0_result = vector.addf %hi0_scaled, %hi0_mma : vector<8xf32> + vector.fragment.store %hi0_result, %cache_view[%c0, %c64] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %hi1_state = vector.fragment.load %upd_tile[%c0, %c16] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %hi1_scaled = vector.mulf %hi1_state, %state_factor_v : vector<8xf32> + %hi1_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi1_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi1_rhs = vector.fragment.load %ka_rhs[%c0, %c80] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi1_mma = vector.mma %hi1_lhs, %hi1_rhs, %hi1_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi1_result = vector.addf %hi1_scaled, %hi1_mma : vector<8xf32> + vector.fragment.store %hi1_result, %cache_view[%c0, %c80] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %hi2_state = vector.fragment.load %upd_tile[%c0, %c32] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %hi2_scaled = vector.mulf %hi2_state, %state_factor_v : vector<8xf32> + %hi2_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi2_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi2_rhs = vector.fragment.load %ka_rhs[%c0, %c96] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi2_mma = vector.mma %hi2_lhs, %hi2_rhs, %hi2_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi2_result = vector.addf %hi2_scaled, %hi2_mma : vector<8xf32> + vector.fragment.store %hi2_result, %cache_view[%c0, %c96] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %hi3_state = vector.fragment.load %upd_tile[%c0, %c48] shape [%c16, %c16] : view<16x68xf32, %lay_upd> -> vector<8xf32> + %hi3_scaled = vector.mulf %hi3_state, %state_factor_v : vector<8xf32> + %hi3_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi3_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi3_rhs = vector.fragment.load %ka_rhs[%c0, %c112] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi3_mma = vector.mma %hi3_lhs, %hi3_rhs, %hi3_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi3_result = vector.addf %hi3_scaled, %hi3_mma : vector<8xf32> + vector.fragment.store %hi3_result, %cache_view[%c0, %c112] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + template.return +} + +template.def<@llm.gated_delta_net.snapshot_fragment_prefix.body> device @llm_gated_delta_net_snapshot_fragment_prefix_body(%state_factor: f32, %cache_base: index, %state_offset: index, %wcol: index, %row0: index, %rowh: index, %wave: index, %lane: index, %lds: buffer, %snapshot_cache: buffer, %lo0_state: vector<8xf32>, %lo1_state: vector<8xf32>, %lo2_state: vector<8xf32>, %lo3_state: vector<8xf32>, %hi0_state: vector<8xf32>, %hi1_state: vector<8xf32>, %hi2_state: vector<8xf32>, %hi3_state: vector<8xf32>, %fz8: vector<8xf32>) { + %c0o = index.constant 0 : offset + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c48 = index.constant 48 : index + %c64 = index.constant 64 : index + %c68 = index.constant 68 : index + %c80 = index.constant 80 : index + %c96 = index.constant 96 : index + %c112 = index.constant 112 : index + %c128 = index.constant 128 : index + %c4352 = index.constant 4352 : index + %c6528 = index.constant 6528 : index + %c11072 = index.constant 11072 : index + %ka_o = index.constant 4352 : offset + %lay_dcol = encoding.layout.strided [24, 1] : encoding + %lay_katok = encoding.layout.strided [1, 24] : encoding + %lay_upd = encoding.layout.strided [68, 1] : encoding + %lay_state = encoding.layout.strided [128, 1] : encoding + %wave_bytes = index.mul %wave, %c6528 : index + %wave_base = index.add %wave_bytes, %c11072 : index + %wave_off = index.cast %wave_base : index to offset + %d_base_i = index.add %wave_base, %c4352 : index + %d_base = index.cast %d_base_i : index to offset + %d_lhs = buffer.view %lds[%d_base] : buffer -> view<16x16xf16, %lay_dcol> + %ka_rhs = buffer.view %lds[%ka_o] : buffer -> view<16x128xf16, %lay_katok> + %upd_flat = buffer.view %lds[%wave_off] : buffer -> view<1088xf32> + %upd_tile = buffer.view %lds[%wave_off] : buffer -> view<16x68xf32, %lay_upd> + %cache_g = buffer.assume.memory_space %snapshot_cache : buffer + %cache_state_base = index.add %cache_base, %state_offset : index + %wave_state = index.mul %wcol, %c128 : index + %cache_wave_base = index.add %cache_state_base, %wave_state : index + %cache_byte_i = index.mul %cache_wave_base, %c4 : index + %cache_byte = index.cast %cache_byte_i : index to offset + %cache_view = buffer.view %cache_g[%cache_byte] : buffer -> view<16x128xf32, %lay_state> + %state_factor_v = vector.splat %state_factor : vector<8xf32> + %lane_low = index.cmp ult, %lane, %c16 : index + %lane_high = index.cmp uge, %lane, %c16 : index + + %lo0_scaled = vector.mulf %lo0_state, %state_factor_v : vector<8xf32> + %lo0_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo0_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo0_rhs = vector.fragment.load %ka_rhs[%c0, %c0] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo0_mma = vector.mma %lo0_lhs, %lo0_rhs, %lo0_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo0_result = vector.addf %lo0_scaled, %lo0_mma : vector<8xf32> + vector.fragment.store %lo0_result, %cache_view[%c0, %c0] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %lo1_scaled = vector.mulf %lo1_state, %state_factor_v : vector<8xf32> + %lo1_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo1_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo1_rhs = vector.fragment.load %ka_rhs[%c0, %c16] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo1_mma = vector.mma %lo1_lhs, %lo1_rhs, %lo1_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo1_result = vector.addf %lo1_scaled, %lo1_mma : vector<8xf32> + vector.fragment.store %lo1_result, %cache_view[%c0, %c16] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %lo2_scaled = vector.mulf %lo2_state, %state_factor_v : vector<8xf32> + %lo2_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo2_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo2_rhs = vector.fragment.load %ka_rhs[%c0, %c32] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo2_mma = vector.mma %lo2_lhs, %lo2_rhs, %lo2_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo2_result = vector.addf %lo2_scaled, %lo2_mma : vector<8xf32> + vector.fragment.store %lo2_result, %cache_view[%c0, %c32] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %lo3_scaled = vector.mulf %lo3_state, %state_factor_v : vector<8xf32> + %lo3_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %lo3_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %lo3_rhs = vector.fragment.load %ka_rhs[%c0, %c48] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %lo3_mma = vector.mma %lo3_lhs, %lo3_rhs, %lo3_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %lo3_result = vector.addf %lo3_scaled, %lo3_mma : vector<8xf32> + vector.fragment.store %lo3_result, %cache_view[%c0, %c48] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + kernel.barrier scope(subgroup) ordering(acq_rel) + %hi0_scaled = vector.mulf %hi0_state, %state_factor_v : vector<8xf32> + %hi0_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi0_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi0_rhs = vector.fragment.load %ka_rhs[%c0, %c64] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi0_mma = vector.mma %hi0_lhs, %hi0_rhs, %hi0_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi0_result = vector.addf %hi0_scaled, %hi0_mma : vector<8xf32> + vector.fragment.store %hi0_result, %cache_view[%c0, %c64] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %hi1_scaled = vector.mulf %hi1_state, %state_factor_v : vector<8xf32> + %hi1_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi1_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi1_rhs = vector.fragment.load %ka_rhs[%c0, %c80] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi1_mma = vector.mma %hi1_lhs, %hi1_rhs, %hi1_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi1_result = vector.addf %hi1_scaled, %hi1_mma : vector<8xf32> + vector.fragment.store %hi1_result, %cache_view[%c0, %c80] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %hi2_scaled = vector.mulf %hi2_state, %state_factor_v : vector<8xf32> + %hi2_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi2_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi2_rhs = vector.fragment.load %ka_rhs[%c0, %c96] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi2_mma = vector.mma %hi2_lhs, %hi2_rhs, %hi2_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi2_result = vector.addf %hi2_scaled, %hi2_mma : vector<8xf32> + vector.fragment.store %hi2_result, %cache_view[%c0, %c96] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + %hi3_scaled = vector.mulf %hi3_state, %state_factor_v : vector<8xf32> + %hi3_init = vector.fragment %fz8 shape [%c16, %c16] : vector<8xf32> + %hi3_lhs = vector.fragment.load %d_lhs[%c0, %c0] shape [%c16, %c16] : view<16x16xf16, %lay_dcol> -> vector<16xf16> + %hi3_rhs = vector.fragment.load %ka_rhs[%c0, %c112] shape [%c16, %c16] : view<16x128xf16, %lay_katok> -> vector<16xf16> + %hi3_mma = vector.mma %hi3_lhs, %hi3_rhs, %hi3_init : vector<16xf16>, vector<16xf16>, vector<8xf32> + %hi3_result = vector.addf %hi3_scaled, %hi3_mma : vector<8xf32> + vector.fragment.store %hi3_result, %cache_view[%c0, %c112] shape [%c16, %c16] : vector<8xf32>, view<16x128xf32, %lay_state> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/llm_attention_qkv_matmul_postops.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/llm_attention_qkv_matmul_postops.loom new file mode 100644 index 000000000000..3d1868a39790 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/llm_attention_qkv_matmul_postops.loom @@ -0,0 +1,101 @@ +template.decl @llm.attention_qkv.normal_rope_vector4(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %values: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.rope_vector4(%token: index, %channel: index, %token_count0: index, %head_size0: index, %head_count0: index, %positions: buffer, %theta: buffer, %freq_factors: buffer, %values: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_cache_vector4(%output_format: index, %token0: index, %channel: index, %token_count0: index, %cache_row_count0: index, %output_size0: index, %indices: buffer, %values: vector<4xf32>, %cache: buffer) + +template.decl @llm.attention_qkv.store_query_vector4(%token0: index, %channel: index, %token_count0: index, %head_size0: index, %head_count0: index, %values: vector<4xf32>, %output: buffer) + +func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +template.def<@llm.attention_qkv.normal_rope_vector4> device @llm_attention_qkv_normal_rope_vector4(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %values: vector<4xf32>) -> (vector<4xf32>) { + %x0 = vector.extract %values[0] : vector<4xf32> -> f32 + %y0 = vector.extract %values[1] : vector<4xf32> -> f32 + %x1 = vector.extract %values[2] : vector<4xf32> -> f32 + %y1 = vector.extract %values[3] : vector<4xf32> -> f32 + %x_values = vector.from_elements %x0, %x1 : vector<2xf32> + %y_values = vector.from_elements %y0, %y1 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta, %freq_factors, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %rotated_x0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 + %rotated_x1 = vector.extract %rotated_x[1] : vector<2xf32> -> f32 + %rotated_y0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 + %rotated_y1 = vector.extract %rotated_y[1] : vector<2xf32> -> f32 + %rotated = vector.from_elements %rotated_x0, %rotated_y0, %rotated_x1, %rotated_y1 : vector<4xf32> + template.return %rotated : vector<4xf32> +} + +template.def<@llm.attention_qkv.rope_vector4> device @llm_attention_qkv_rope_vector4(%token: index, %channel: index, %token_count0: index, %head_size0: index, %head_count0: index, %positions: buffer, %theta: buffer, %freq_factors: buffer, %values: vector<4xf32>) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %head_size = index.assume %head_size0 [range(%head_size0, 4, 1024), mul(%head_size0, 4)] : index + %head_count = index.assume %head_count0 [range(%head_count0, 1, 64)] : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c0_offset = index.constant 0 : offset + %half_head_size = index.div %head_size, %c2 : index + %head_channel = index.rem %channel, %head_size : index + %channel_packet = index.div %head_channel, %c4 : index + %theta_channel = index.mul %channel_packet, %c2 : index + %positions_noalias, %theta_noalias, %freq_factors_noalias = buffer.assume.noalias %positions, %theta, %freq_factors : buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> vector<2xf32> + %rotated = template.apply<@llm.attention_qkv.normal_rope_vector4>(%position, %theta_packet, %freq_factors_packet, %values) : (f32, vector<2xf32>, vector<2xf32>, vector<4xf32>) -> (vector<4xf32>) + template.return %rotated : vector<4xf32> +} + +template.def<@llm.attention_qkv.store_query_vector4> device @llm_attention_qkv_store_query_vector4(%token0: index, %channel: index, %token_count0: index, %head_size0: index, %head_count0: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %head_size = index.assume %head_size0 [range(%head_size0, 4, 1024), mul(%head_size0, 4)] : index + %head_count = index.assume %head_count0 [range(%head_count0, 1, 64)] : index + %c0_offset = index.constant 0 : offset + %head = index.div %channel, %head_size : index + %head_channel = index.rem %channel, %head_size : index + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%head_count]x[%head_size]xf32> + vector.store %values, %output_view[%token, %head, %head_channel] : vector<4xf32>, view<[%token_count]x[%head_count]x[%head_size]xf32> + template.return +} + +template.def<@llm.attention_qkv.store_cache_vector4> device @llm_attention_qkv_store_cache_vector4(%output_format: index, %token0: index, %channel: index, %token_count0: index, %cache_row_count0: index, %output_size0: index, %indices: buffer, %values: vector<4xf32>, %cache: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %is_f16 = index.cmp eq, %output_format, %c16 : index + %is_f32 = index.cmp eq, %output_format, %c32 : index + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %indices_noalias, %cache_noalias = buffer.assume.noalias %indices, %cache : buffer, buffer + %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + %publish = scalar.andi %valid_index, %valid_row : i1 + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<4xf16>, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<4xf32>, view<[%cache_row_count]x[%output_size]xf32> + } + template.return +} 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/q4_q6_q8_1_x4_qkv.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom new file mode 100644 index 000000000000..63f8c97a55a9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom @@ -0,0 +1,113 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.attention.qkv_quantized.body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +func.decl @qwen3_moe_q4k_q8_1_x4_paired_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_q6k_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) + +// Device body shared by the ordinary projection and completion-fused +// postprocess exports. Each caller owns its boundary after the raw row stores. +template.def<@qwen3_moe.attention.qkv_quantized.body> device @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %query_size, %key_value_size = index.assume %query_size0, %key_value_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), mul(%query_size0, %key_value_size0)] : index, index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %global_channel = index.add %channel_base, %subgroup : index + %key_value_end = index.add %query_size, %key_value_size : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %total_output_size = index.add %query_size, %key_value_output_size : index + %valid_channel = index.cmp ult, %global_channel, %total_output_size : index + %is_query = index.cmp ult, %global_channel, %query_size : index + %is_key = index.cmp ult, %global_channel, %key_value_end : index + %key_value_channel = index.rem %global_channel, %key_value_size : index + %value_uses_q6 = index.cmp eq, %value_uses_q6_index, %c1 : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %quant_block_count = index.div %hidden_size, %c256 : index + %q4_row_bytes = index.scale %quant_block_count, %c144_bytes : index, offset -> offset + %q6_row_bytes = index.scale %quant_block_count, %c210_bytes : index, offset -> offset + %q8_group_count = index.div %hidden_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %query_weight_noalias, %key_weight_noalias, %value_weight_noalias, %query_output_noalias, %key_output_noalias, %value_output_noalias = buffer.assume.noalias %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %projection_sum = scf.if %is_query -> (f32) { + %row_byte_base = index.scale %global_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %query_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %key_or_value_sum = scf.if %is_key -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %key_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %value_sum = scf.if %value_uses_q6 -> (f32) { + %row_byte_base = index.scale %key_value_channel, %q6_row_bytes : index, offset -> offset + %sum = func.call @ggml_q6k_q8_1_x4_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } + scf.yield %value_sum : f32 + } + scf.yield %key_or_value_sum : f32 + } + scf.yield %projection_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_size]xf32> + %key_output_view = buffer.view %key_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + %value_output_view = buffer.view %value_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + scf.if %is_query { + view.store %dot, %query_output_view[%token, %global_channel] : f32, view<[%launch_token_count]x[%query_size]xf32> + } else { + scf.if %is_key { + view.store %dot, %key_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } else { + view.store %dot, %value_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> + } + } + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom new file mode 100644 index 000000000000..7534dcbe7b3c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom @@ -0,0 +1,200 @@ +// Packed Q4_K x Q8_1 x4 row contraction shared by dense wrappers. +// The symbols and bodies are preserved from the retired routed gate-up packet. + +func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : 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 + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +func.def inline @qwen3_moe_q4k_q8_1_x4_paired_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, 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 + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %padded_input_size = index.add %bounded_input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %block_lane0 = index.rem %bounded_lane, %c8 : index + %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index + %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %contribution = scf.if %valid_q4_block -> (f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + scf.yield %pair : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %iteration_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom new file mode 100644 index 000000000000..4acae68782f2 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom @@ -0,0 +1,75 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.dense.q4k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +config.decl @qwen3_moe.dense_quantized.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @qwen3_moe.dense_quantized.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.dense_quantized.output_accumulation : %value: index where [range(%value, 0, 1)] + +func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) + +template.def<@qwen3_moe.dense.q4k_q8_1_x4.body> device @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @qwen3_moe.dense_quantized.input_size : index + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144_bytes = index.constant 144 : offset + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q4_block_count, %c144_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer + %lane_sum = scf.if %publish_output -> (f32) { + %channel_sum = scf.if %valid_channel -> (f32) { + %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + scf.yield %channel_sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_sum : f32 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_output { + scf.if %valid_channel { + scf.if %is_lane_zero { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%bounded_output_size]xf32> + %result = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %channel] : view<[%launch_token_count]x[%bounded_output_size]xf32> -> f32 + %sum = scalar.addf %residual, %dot : f32 + scf.yield %sum : f32 + } else { + scf.yield %dot : f32 + } + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%bounded_output_size]xf32> + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_f32_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_f32_decode.loom new file mode 100644 index 000000000000..b5b9f33be8e5 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_f32_decode.loom @@ -0,0 +1,213 @@ +// Common direct Q6_K/F32 decode helpers used by dense and routed vector kernels. +// Computes an ordered four-element F32 dot product. The scalar recurrence +// mirrors the GLSL oracle and gives target scheduling four independent dot +// chains per decoded Q6_K block. +func.def inline @ggml_q6k_dot4_f32(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (f32) { + %c0 = scalar.constant 0.0 : f32 + %lhs0 = vector.extract %lhs[0] : vector<4xf32> -> f32 + %lhs1 = vector.extract %lhs[1] : vector<4xf32> -> f32 + %lhs2 = vector.extract %lhs[2] : vector<4xf32> -> f32 + %lhs3 = vector.extract %lhs[3] : vector<4xf32> -> f32 + %rhs0 = vector.extract %rhs[0] : vector<4xf32> -> f32 + %rhs1 = vector.extract %rhs[1] : vector<4xf32> -> f32 + %rhs2 = vector.extract %rhs[2] : vector<4xf32> -> f32 + %rhs3 = vector.extract %rhs[3] : vector<4xf32> -> f32 + %sum0 = scalar.fmaf %lhs0, %rhs0, %c0 : f32 + %sum1 = scalar.fmaf %lhs1, %rhs1, %sum0 : f32 + %sum2 = scalar.fmaf %lhs2, %rhs2, %sum1 : f32 + %sum3 = scalar.fmaf %lhs3, %rhs3, %sum2 : f32 + func.return %sum3 : f32 +} + +// Loads one row contraction's activation packet. The two output rows read the +// same packet, but each load remains adjacent to its consumer. Hoisting the +// packet across the intervening LDS barrier halves issued activation traffic +// while extending sixteen F32 live values across synchronization; on gfx1151 +// that longer live range is slower than the repeated cache-hot read. +func.def inline @ggml_q6k_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 + %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 + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %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 + %input0 = vector.load %input_view[%token, %safe_block, %input_index0] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input1 = vector.load %input_view[%token, %safe_block, %input_index1] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input2 = vector.load %input_view[%token, %safe_block, %input_index2] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + %input3 = vector.load %input_view[%token, %safe_block, %input_index3] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> + func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} + +// Loads the signed group scale owned by one lane of a Q6_K block. +func.def inline @ggml_q6k_load_f32_scale(%input_size: index, %row: index, %block: index, %lane: index, %weight: buffer) -> (f32) { + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %c192 = index.constant 192 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.if %valid_block -> (index) { + scf.yield %block : index + } else { + scf.yield %c0 : index + } + %itid = index.rem %bounded_lane, %c16 : index + %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset + %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %weight_block_byte_base, %c192 : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %scale_i8 = view.load %scale_view[%itid] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + func.return %scale : f32 +} + +// Stages one row's group scales and synchronizes before contraction. +func.def inline @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) { + %c0_offset = index.constant 0 : offset + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> + %scale = func.call @ggml_q6k_load_f32_scale(%input_size, %row, %block, %bounded_lane, %weight) : (index, index, index, index, buffer) -> (f32) + view.store %scale, %scale_stage_view[%frame, %bounded_lane] : f32, view<[%frame_count]x64xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +// Computes one lane contribution for one output row and one four-block +// cohort after its group scales have been staged. +func.def inline @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c128 = index.constant 128 : offset + %c208 = index.constant 208 : offset + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high0_mask = vector.constant 50529027 : vector<1xi32> + %high2_mask = vector.constant 202116108 : vector<1xi32> + %high4_mask = vector.constant 808464432 : vector<1xi32> + %high6_mask = vector.constant -1061109568 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %block_count = index.div %input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.if %valid_block -> (index) { + scf.yield %block : index + } else { + scf.yield %c0 : index + } + %itid = index.rem %bounded_lane, %c16 : index + %cohort = index.div %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %vector_quarter = index.div %vector_index, %c4 : index + %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset + %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %weight_block_byte_base, %c128 : offset + %d_byte_base = index.add %weight_block_byte_base, %c208 : offset + %ql_view = buffer.view %weight[%weight_block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> + %scale_cohort_base = index.mul %cohort, %c16 : index + %scale_half_base = index.mul %vector_half, %c8 : index + %scale_index0 = index.add %scale_half_base, %vector_quarter : index + %scale_index1 = index.add %scale_index0, %c2 : index + %scale_index2 = index.add %scale_index0, %c4 : index + %scale_index3 = index.add %scale_index0, %c6 : index + %stage_scale_index0 = index.add %scale_cohort_base, %scale_index0 : index + %stage_scale_index1 = index.add %scale_cohort_base, %scale_index1 : index + %stage_scale_index2 = index.add %scale_cohort_base, %scale_index2 : index + %stage_scale_index3 = index.add %scale_cohort_base, %scale_index3 : index + %scale0 = view.load %scale_stage_view[%frame, %stage_scale_index0] : view<[%frame_count]x64xf32> -> f32 + %scale1 = view.load %scale_stage_view[%frame, %stage_scale_index1] : view<[%frame_count]x64xf32> -> f32 + %scale2 = view.load %scale_stage_view[%frame, %stage_scale_index2] : view<[%frame_count]x64xf32> -> f32 + %scale3 = view.load %scale_stage_view[%frame, %stage_scale_index3] : view<[%frame_count]x64xf32> -> f32 + %ql_half_word_base = index.mul %vector_half, %c16 : index + %ql_word_index00 = index.add %ql_half_word_base, %vector_index : index + %ql_word_index10 = index.add %ql_word_index00, %c8 : index + %qh_half_word_base = index.mul %vector_half, %c8 : index + %qh_word_index0 = index.add %qh_half_word_base, %vector_index : index + %ql_word_index0, %ql_word_index1, %qh_word_index = index.assume %ql_word_index00, %ql_word_index10, %qh_word_index0 [range(%ql_word_index00, 0, 23), range(%ql_word_index10, 8, 31), range(%qh_word_index0, 0, 15)] : index, index, index + %ql_word0 = vector.load %ql_view[%ql_word_index0] : view<32xi32> -> vector<1xi32> + %ql_word1 = vector.load %ql_view[%ql_word_index1] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql0 = vector.andi %ql_word0, %nibble_mask : vector<1xi32> + %ql1 = vector.andi %ql_word1, %nibble_mask : vector<1xi32> + %ql_word0_high = vector.shrui %ql_word0, %c4_i32v : vector<1xi32> + %ql_word1_high = vector.shrui %ql_word1, %c4_i32v : vector<1xi32> + %ql2 = vector.andi %ql_word0_high, %nibble_mask : vector<1xi32> + %ql3 = vector.andi %ql_word1_high, %nibble_mask : vector<1xi32> + %qh0_low = vector.andi %qh_word, %high0_mask : vector<1xi32> + %qh1_low = vector.andi %qh_word, %high2_mask : vector<1xi32> + %qh2 = vector.andi %qh_word, %high4_mask : vector<1xi32> + %qh3_high = vector.andi %qh_word, %high6_mask : vector<1xi32> + %qh0 = vector.shli %qh0_low, %c4_i32v : vector<1xi32> + %qh1 = vector.shli %qh1_low, %c2_i32v : vector<1xi32> + %qh3 = vector.shrui %qh3_high, %c2_i32v : vector<1xi32> + %code0 = vector.ori %ql0, %qh0 : vector<1xi32> + %code1 = vector.ori %ql1, %qh1 : vector<1xi32> + %code2 = vector.ori %ql2, %qh2 : vector<1xi32> + %code3 = vector.ori %ql3, %qh3 : vector<1xi32> + %code0_i8 = vector.bitcast %code0 : vector<1xi32> to vector<4xi8> + %code1_i8 = vector.bitcast %code1 : vector<1xi32> to vector<4xi8> + %code2_i8 = vector.bitcast %code2 : vector<1xi32> to vector<4xi8> + %code3_i8 = vector.bitcast %code3 : vector<1xi32> to vector<4xi8> + %code0_f32 = vector.uitofp %code0_i8 : vector<4xi8> to vector<4xf32> + %code1_f32 = vector.uitofp %code1_i8 : vector<4xi8> to vector<4xf32> + %code2_f32 = vector.uitofp %code2_i8 : vector<4xi8> to vector<4xf32> + %code3_f32 = vector.uitofp %code3_i8 : vector<4xi8> to vector<4xf32> + %q0 = vector.subf %code0_f32, %c32_f32v : vector<4xf32> + %q1 = vector.subf %code1_f32, %c32_f32v : vector<4xf32> + %q2 = vector.subf %code2_f32, %c32_f32v : vector<4xf32> + %q3 = vector.subf %code3_f32, %c32_f32v : vector<4xf32> + %dot0 = func.call @ggml_q6k_dot4_f32(%input0, %q0) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot1 = func.call @ggml_q6k_dot4_f32(%input1, %q1) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot2 = func.call @ggml_q6k_dot4_f32(%input2, %q2) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot3 = func.call @ggml_q6k_dot4_f32(%input3, %q3) : (vector<4xf32>, vector<4xf32>) -> (f32) + %scaled3 = scalar.mulf %dot3, %scale3 : f32 + %scaled2 = scalar.fmaf %dot2, %scale2, %scaled3 : f32 + %scaled1 = scalar.fmaf %dot1, %scale1, %scaled2 : f32 + %scaled0 = scalar.fmaf %dot0, %scale0, %scaled1 : f32 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %contribution0 = scalar.mulf %scaled0, %d : f32 + %contribution = scf.if %valid_block -> (f32) { + scf.yield %contribution0 : f32 + } else { + scf.yield %c0_f32 : f32 + } + func.return %contribution : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom new file mode 100644 index 000000000000..49df978e7e09 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom @@ -0,0 +1,204 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.linear_q6k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) + +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_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) + +func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { + %c1_i32v = vector.constant 1 : vector<1xi32> + %c2_i32v = vector.constant 2 : vector<1xi32> + %low5_mask = vector.constant 522133279 : vector<1xi32> + %bit5_mask = vector.constant 538976288 : vector<1xi32> + %sign_mask = vector.constant -522133280 : vector<1xi32> + %low5 = vector.andi %code, %low5_mask : vector<1xi32> + %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> + %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> + %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> + %high01 = vector.ori %bit5, %bit6 : vector<1xi32> + %high = vector.ori %high01, %bit7 : vector<1xi32> + %sign = vector.xori %high, %sign_mask : vector<1xi32> + %signed_i32 = vector.ori %low5, %sign : vector<1xi32> + %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> + func.return %signed : vector<4xi8> +} + +func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %c0_i32v = vector.constant 0 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 808464432 : vector<1xi32> + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %lane_mod8 = index.rem %bounded_lane, %c8 : index + %lane_mod16 = index.rem %bounded_lane, %c16 : index + %lane_div16 = index.div %bounded_lane, %c16 : index + %lane_div8_in_16 = index.div %lane_mod16, %c8 : index + %lane_div4_in_16 = index.div %lane_mod16, %c4 : index + %qh_high_base = index.mul %lane_div16, %c8 : index + %qh_index0 = index.add %qh_high_base, %lane_mod8 : index + %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index + %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> + %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index + %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 + %q8_block_base = index.mul %q6_block, %c8 : index + %q8_high_add = index.mul %lane_div16, %c4 : index + %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index + %scale_high_base = index.mul %lane_div16, %c8 : index + %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { + %bounded_part = index.assume %part [range(%part, 0, 1)] : index + %part_shift_index = index.mul %bounded_part, %c4 : index + %part_shift_i32 = index.cast %part_shift_index : index to i32 + %part_shift = vector.splat %part_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> + %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) + %q8_part_add = index.mul %bounded_part, %c2 : index + %q8_block_part = index.add %q8_block_base, %q8_part_add : index + %q8_block = index.add %q8_block_part, %q8_quadrant : index + %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) + %scale_lane = index.add %scale_lane0, %part_shift_index : index + %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 + %scale = scalar.sitofp %scale_i8 : i8 to 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 + %scaled0 = scalar.mulf %dot_f32, %scale : f32 + %scaled1 = scalar.mulf %scaled0, %d : f32 + %contribution = scalar.mulf %scaled1, %q8_d : f32 + %next = scalar.addf %accumulator, %contribution : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} + +// Computes one lane's partial for a complete Q6_K row. Callers choose how +// rows and activations are tiled, then reduce the returned value by subgroup. +func.def inline @ggml_q6k_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, 256, 32768), mul(%input_size, 256)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { + %contribution = func.call @ggml_q6k_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 %block_acc, %contribution : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Maps four Q6_K blocks across the four 16-lane partitions of a wave64. Each +// physical lane contracts the pair of virtual wave32 lanes that own the same +// packed positions in the low and high halves of one block. The subgroup +// reduction therefore combines four complete blocks per loop iteration while +// preserving the canonical wave32 unpacking and dot-product primitive. +func.def inline @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%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, 256, 32768), mul(%input_size, 256)] : index + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %block_count = index.div %bounded_input_size, %c256 : index + %padded_block_count = index.add %block_count, %c3 : index + %iteration_count = index.div %padded_block_count, %c4 : index + %block_in_iteration = index.div %bounded_lane, %c16 : index + %virtual_lane0 = index.rem %bounded_lane, %c16 : index + %virtual_lane1 = index.add %virtual_lane0, %c16 : index + %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%row_acc = %c0_f32 : f32) -> (f32) { + %block_base = index.mul %iteration, %c4 : index + %block0 = index.add %block_base, %block_in_iteration : index + %valid_block = index.cmp ult, %block0, %block_count : index + %block_acc = scf.if %valid_block -> (f32) { + %low = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane0) : (buffer, offset, buffer, offset, index, index) -> (f32) + %high = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane1) : (buffer, offset, buffer, offset, index, index) -> (f32) + %sum = scalar.addf %low, %high : f32 + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %row_acc, %block_acc : f32 + scf.yield %next : f32 + } + func.return %result : f32 +} + +// Contracts one output channel for the calling Q6_K projection kernel. The +// caller owns publication so the same canonical contraction can feed either a +// dense logits tensor or an endpoint reduction. +template.def<@ggml.linear_q6k_q8_1_x4.body> device @ggml_linear_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c210_bytes = index.constant 210 : offset + %c256 = index.constant 256 : index + %c144_bytes = index.constant 144 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset + %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset + %lane_acc = scf.if %publish_output -> (f32) { + %channel_acc = scf.if %valid_channel -> (f32) { + %value = func.call @ggml_q6k_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) + scf.yield %value : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + scf.yield %channel_acc : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %dot = kernel.subgroup.reduce %lane_acc : f32 + template.return %dot, %token, %channel, %valid_channel, %launch_token_count, %bounded_output_size : f32, index, index, i1, index, index +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..c7bf51776f16 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,1459 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.mul_mat.activation_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Shared contiguous 2D F32 activation skinny matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4 +// ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def pure inline @ggml_q8_lowtoken_geometry(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %q4_format = index.constant 44 : index + %q6_format = index.constant 46 : index + %q8_format = index.constant 9 : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %native_q4 = index.cmp eq, %weight_format, %q4_format : index + %native_q6 = index.cmp eq, %weight_format, %q6_format : index + %q8_input = index.cmp eq, %activation_format, %q8_format : index + %q4_min_blocks = index.constant 327680 : index + %blocks = index.div %input_size, %c256 : index + %weight_blocks = index.mul %blocks, %output_size : index + %large_payload = index.cmp uge, %weight_blocks, %q4_min_blocks : index + %q4_can_stage = scalar.andi %native_q4, %large_payload : i1 + %native_staged = scalar.ori %q4_can_stage, %native_q6 : i1 + %native_q8 = scalar.andi %native_staged, %q8_input : i1 + %q4_min_k = index.constant 16384 : index + %q6_min_k = index.constant 8192 : index + %min_k = scf.select %native_q4, %q4_min_k, %q6_min_k : index + %min_n = index.constant 4096 : index + %five_tokens = index.cmp eq, %token_capacity, %c5 : index + %long_k = index.cmp uge, %input_size, %min_k : index + %wide_n = index.cmp uge, %output_size, %min_n : index + %contracts = index.cmp ule, %output_size, %input_size : index + %large = scalar.andi %long_k, %wide_n : i1 + %long_reduction = scalar.andi %large, %contracts : i1 + %staged_shape = scalar.andi %five_tokens, %long_reduction : i1 + %staged = scalar.andi %native_q8, %staged_shape : i1 + %rows = scf.select %native_q4, %c64, %c32 : index + %result = scf.select %staged, %rows, %c0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %block_bytes = index.constant 144 : index + %half_cache_bytes = index.constant 16777216 : index + %cache_bytes = index.constant 33554432 : index + %payload_bytes = index.mul %weight_blocks, %block_bytes : index + %large_half = index.cmp uge, %payload_bytes, %half_cache_bytes : index + %fits_cache = index.cmp ule, %payload_bytes, %cache_bytes : index + %cache_range = scalar.andi %large_half, %fits_cache : i1 + %remainder = index.rem %blocks, %c2 : index + %even_blocks = index.cmp eq, %remainder, %c0 : index + %reduces_rows = index.cmp ult, %output_size, %input_size : index + %native_q4_q8 = scalar.andi %native_q4, %q8_input : i1 + %plain = index.cmp eq, %result, %c0 : index + %q4_plain = scalar.andi %native_q4_q8, %plain : i1 + %five_wide = scalar.andi %five_tokens, %wide_n : i1 + %reduction = scalar.andi %reduces_rows, %even_blocks : i1 + %partition_shape = scalar.andi %five_wide, %reduction : i1 + %partition_payload = scalar.andi %partition_shape, %cache_range : i1 + %partitioned = scalar.andi %q4_plain, %partition_payload : i1 + %k_partitions = scf.select %partitioned, %c2, %c1 : index + func.return %result, %k_partitions : index, index +} + +func.def pure inline @ggml_mul_mat_uses_lowtoken_dot(%token_capacity: index, %input_size: index) -> (i1) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c256 = index.constant 256 : index + %few_tokens = index.cmp ule, %token_capacity, %c5 : index + %tail = index.rem %input_size, %c256 : index + %whole_blocks = index.cmp eq, %tail, %c0 : index + %use_dot = scalar.andi %few_tokens, %whole_blocks : i1 + func.return %use_dot : i1 +} + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch> pure @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c5 = index.constant 5 : index + %c8 = index.constant 8 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %small_output = index.cmp ult, %output_size, %c64 : index + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%token_capacity, %input_size, %output_size, %weight_format) : (index, index, index, index) -> (index, index) + %staged = index.cmp ugt, %staged_rows, %c0 : index + %staged_rounding = index.sub %staged_rows, %c1 : index + %staged_threads = index.mul %staged_rows, %c16 : index + %cohort_rows = index.div %c8, %k_partitions : index + %cohort_rounding = index.sub %cohort_rows, %c1 : index + %base_row_tile = scf.select %small_output, %c4, %cohort_rows : index + %c1024 = index.constant 1024 : index + %c4096 = index.constant 4096 : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %batched_tokens = index.cmp ugt, %token_capacity, %c1 : index + %long_reduction = index.cmp uge, %input_size, %c4096 : index + %partition_tail = index.rem %input_size, %c1024 : index + %full_partitions = index.cmp eq, %partition_tail, %c0 : index + %batch_f32 = scalar.andi %batched_tokens, %f32_input : i1 + %partitionable = scalar.andi %long_reduction, %full_partitions : i1 + %partition_rows = scalar.andi %batch_f32, %partitionable : i1 + %narrow_threads = scf.select %partition_rows, %c1024, %c256 : index + %base_workgroup_size = scf.select %small_output, %narrow_threads, %c128 : index + %base_rounding = scf.select %small_output, %c3, %cohort_rounding : index + %row_rounding = scf.select %staged, %staged_rounding, %base_rounding : index + %row_tile = scf.select %staged, %staged_rows, %base_row_tile : index + %workgroup_size = scf.select %staged, %staged_threads, %base_workgroup_size : index + %rounded_n = index.add %output_size, %row_rounding : index + %direct_tiles = index.div %rounded_n, %row_tile : index + template.return %direct_tiles, %c1, %c1, %workgroup_size : index, index, index, index +} + +func.def inline @ggml_lowtoken_combine_k_partitions(%value: f32, %partitions: index) -> (f32) { + %c1 = index.constant 1 : index + %partitioned = index.cmp ugt, %partitions, %c1 : index + %result = scf.if %partitioned -> (f32) { + %offset = scalar.constant 16 : i32 + %width = scalar.constant 32 : i32 + %peer, %ok = kernel.subgroup.shuffle %value, %offset, %width : f32, i32, i32 + %sum = scalar.addf %value, %peer : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_lowtoken_reduce_cohort_f32(%value: f32) -> (f32) { + %width = scalar.constant 32 : i32 + %off1 = scalar.constant 1 : i32 + %peer1, %ok1 = kernel.subgroup.shuffle %value, %off1, %width : f32, i32, i32 + %sum1 = scalar.addf %value, %peer1 : f32 + %off2 = scalar.constant 2 : i32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %off2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %off4 = scalar.constant 4 : i32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %off4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %off8 = scalar.constant 8 : i32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %off8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + func.return %sum8 : f32 +} + +// Narrow low-token projections assign a full wave to each row. +// Both cohort sizes use the existing wrapper publication hooks. +func.def inline @ggml_lowtoken_reduce_f32(%values: vector<4xf32>, %full_wave: i1) -> (f32) { + %zero = scalar.constant 0.0 : f32 + %width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %sum = vector.reduce %values, %zero : vector<4xf32>, f32 + %peer1, %ok1 = kernel.subgroup.shuffle %sum, %xor1, %width : f32, i32, i32 + %sum1 = scalar.addf %sum, %peer1 : f32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %xor2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %xor4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %xor8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + %result = scf.if %full_wave -> (f32) { + %xor16 = scalar.constant 16 : i32 + %peer16, %ok16 = kernel.subgroup.shuffle %sum8, %xor16, %width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer16 : f32 + %lane32 = scalar.constant 32 : i32 + %peer32 = kernel.subgroup.broadcast %sum16 from %lane32 : f32, i32 + %sum64 = scalar.addf %sum16, %peer32 : f32 + scf.yield %sum64 : f32 + } else { + scf.yield %sum8 : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_mul_mat_f32_f32_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %tokens = index.assume %token_count [range(%token_count, 2, 5)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workgroup_size = kernel.workgroup.size : index + %one_wave_per_row = index.cmp eq, %workgroup_size, %c256 : index + %c1024 = index.constant 1024 : index + %partition_group = index.cmp eq, %workgroup_size, %c1024 : index + %narrow_output = index.cmp ult, %n, %c64 : index + %partitioned = scalar.andi %partition_group, %narrow_output : i1 + %full_wave = scalar.ori %one_wave_per_row, %partitioned : i1 + %k_partitions = scf.select %partitioned, %c4, %c1 : index + %cohort_width = scf.select %full_wave, %c64, %c16 : index + %tile_rows = scf.select %full_wave, %c4, %c8 : index + %wave_rows = scf.select %full_wave, %c1, %c4 : index + %step_width = scf.select %full_wave, %c256, %c64 : index + %steps_per_block = scf.select %full_wave, %c1, %c4 : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %lane16 = index.rem %lane, %cohort_width : index + %cohort = index.div %lane, %cohort_width : index + %tile_base = index.mul %channel_tile, %tile_rows : index + %row_subgroup = index.div %subgroup, %k_partitions : index + %k_partition = index.rem %subgroup, %k_partitions : index + %wave_add = index.mul %row_subgroup, %wave_rows : index + %wave_base = index.add %tile_base, %wave_add : index + %row = index.add %wave_base, %cohort : index + %valid_row = index.cmp ult, %row, %n : index + %packet = index.rem %lane16, %c8 : index + %lane_group = index.div %lane16, %c8 : index + %lane_k = index.mul %lane16, %c4 : index + %blocks = index.div %k, %c256 : index + %steps = index.div %k, %step_width : index + %partition_steps = index.div %steps, %k_partitions : index + %step_begin = index.mul %k_partition, %partition_steps : index + %step_end = index.add %step_begin, %partition_steps : index + %weight_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) + %row_bytes = index.scale %blocks, %weight_bytes : index, offset -> offset + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %a_na, %w_na = buffer.assume.noalias %input, %weight : buffer, buffer + %av = buffer.view %a_na[%base] : buffer -> view<[%tokens]x[%k]xf32> + %total0, %total1, %total2, %total3, %total4 = scf.for %step = [%step_begin to %step_end step %c1](%acc0 = %zero : vector<4xf32>, %acc1 = %zero : vector<4xf32>, %acc2 = %zero : vector<4xf32>, %acc3 = %zero : vector<4xf32>, %acc4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block = index.div %step, %steps_per_block : index + %quarter = index.rem %step, %steps_per_block : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %step_k = index.mul %step, %step_width : index + %kk = index.add %step_k, %lane_k : index + %ww = scf.if %valid_row -> (vector<4xf32>) { + %format44 = index.constant 44 : index + %packed = index.cmp eq, %weight_format, %format44 : index + %wh = scf.if %packed -> (vector<4xf32>) { + %v = func.call @ggml_q4k_native_row64_f32_vector4(%w_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %v = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %w_na, %row_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } + scf.yield %wh : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next0 = scf.if %valid0 -> (vector<4xf32>) { + %a = vector.load %av[%t0, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc0 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc0 : vector<4xf32> + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next1 = scf.if %valid1 -> (vector<4xf32>) { + %a = vector.load %av[%t1, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc1 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc1 : vector<4xf32> + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next2 = scf.if %valid2 -> (vector<4xf32>) { + %a = vector.load %av[%t2, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc2 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc2 : vector<4xf32> + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next3 = scf.if %valid3 -> (vector<4xf32>) { + %a = vector.load %av[%t3, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc3 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc3 : vector<4xf32> + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next4 = scf.if %valid4 -> (vector<4xf32>) { + %a = vector.load %av[%t4, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc4 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc4 : vector<4xf32> + } + scf.yield %next0, %next1, %next2, %next3, %next4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %leader = index.cmp eq, %lane16, %c0 : index + %c5 = index.constant 5 : index + %stage_width = scf.select %partitioned, %c16, %c8 : index + %stage_elements = index.mul %stage_width, %c5 : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%stage_width]xf32> + %local_row0 = index.add %wave_add, %cohort : index + %partial_offset = index.mul %k_partition, %c4 : index + %local_row = index.add %local_row0, %partial_offset : index + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total0, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total1, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total2, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total3, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total4, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %tile_rows, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %tile_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, 4), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %stage_width : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_origin = scf.if %partitioned -> (index) { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 64), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } else { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 36), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values0 = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %values = scf.if %partitioned -> (vector<4xf32>) { + %stage_origin1 = index.add %stage_origin, %c4 : index + %stage_origin2 = index.add %stage_origin1, %c4 : index + %stage_origin3 = index.add %stage_origin2, %c4 : index + %values1 = vector.load %flat_stage[%stage_origin1] : view<[%stage_elements]xf32> -> vector<4xf32> + %values2 = vector.load %flat_stage[%stage_origin2] : view<[%stage_elements]xf32> -> vector<4xf32> + %values3 = vector.load %flat_stage[%stage_origin3] : view<[%stage_elements]xf32> -> vector<4xf32> + %sum01 = vector.addf %values0, %values1 : vector<4xf32> + %sum012 = vector.addf %sum01, %values2 : vector<4xf32> + %sum = vector.addf %sum012, %values3 : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values0 : vector<4xf32> + } + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %epilogue, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %row_rounding = index.sub %tile_rows, %c1 : index + %rounded_n = index.add %n, %row_rounding : index + %output_tiles = index.div %rounded_n, %tile_rows : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile>(%tokens, %n, %output_tiles, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Copy the original Q8_1_x4 records for up to 2048 reduction elements. +// Call uniformly at a slab boundary; only the final slab may be partial. +func.def inline @ggml_q8_1_x4_stage_lowtoken_slab(%input: buffer, %stage: buffer, %capacity: index, %slab_begin: index, %blocks: index, %global_row_bytes: offset) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c144 = index.constant 144 : index + %c576 = index.constant 576 : index + %c18 = index.constant 18 : index + %base = index.constant 0 : offset + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %workgroup_size = kernel.workgroup.size : index + %thread = kernel.workitem.id : index + %stage_packets = index.mul %capacity, %c144 : index + %stage_elements = index.mul %capacity, %c576 : index + %stage_view = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xi32> + %slab_limit = index.add %slab_begin, %c8 : index + %full_slab = index.cmp ule, %slab_limit, %blocks : index + %slab_end = scf.select %full_slab, %slab_limit, %blocks : index + %active_blocks0 = index.sub %slab_end, %slab_begin : index + %active_blocks = index.assume %active_blocks0 [range(%active_blocks0, 1, 8)] : index + %active_packets = index.mul %active_blocks, %c18 : index + %slab_byte_add = index.scale %slab_begin, %block_bytes : index, offset -> offset + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %copy_packet = [%thread to %stage_packets step %workgroup_size] { + %copy_token = index.div %copy_packet, %c144 : index + %copy_inner = index.rem %copy_packet, %c144 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + scf.if %copy_valid { + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %copy_bytes = index.scale %copy_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slab_byte_add : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + %copy_origin = index.mul %copy_packet, %c4 : index + vector.store %values, %stage_view[%copy_origin] : vector<4xi32>, view<[%stage_elements]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_mul_mat_q8_lowtoken_publish(%tokens: index, %n: index, %row_tile: index, %row_base: index, %local_row: index, %block_lane: index, %channel_tile: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %total_g0: f32, %total_g1: f32, %total_g2: f32, %total_g3: f32, %total_g4: f32, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %leader = index.cmp eq, %block_lane, %c0 : index + %width = scalar.constant 32 : i32 + %stage_rows = index.constant 5 : index + %stage_elements = index.mul %stage_rows, %row_tile : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%row_tile]xf32> + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum = scalar.addf %total_g0, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum = scalar.addf %total_g1, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum = scalar.addf %total_g2, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum = scalar.addf %total_g3, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum = scalar.addf %total_g4, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %row_tile, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %row_limit = index.sub %row_tile, %c4 : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, %row_limit), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %row_tile : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_limit = index.sub %stage_elements, %c4 : index + %stage_origin = index.assume %stage_origin0 [range(%stage_origin0, 0, %stage_limit), mul(%stage_origin0, 4)] : index + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %epilogue, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %rounding = index.sub %row_tile, %c1 : index + %rounded_n = index.add %n, %rounding : index + %output_tiles = index.div %rounded_n, %row_tile : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile>(%tokens, %n, %output_tiles, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_q4_prefetch_native_words(%weight: buffer, %row_group_block: index, %block: index, %row_lane: index, %payload_field: index, %payload_word: index) -> (vector<4xi32>, vector<2xi32>) { + %zero = index.constant 0 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %group_block = index.add %row_group_block, %block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %header_view = buffer.view %weight[%wb] : buffer -> view<4xi32> + %code_view = buffer.view %weight[%wc] : buffer -> view<8x64x4xi32> + %headers = vector.load %header_view[%zero] : view<4xi32> -> vector<4xi32> + %raw = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + func.return %headers, %raw : vector<4xi32>, vector<2xi32> +} + +// Native Q4_K bytes and Q8_1_x4 activations, with one 16-lane cohort per row. +// Retain the wrapper publication hooks so bias, residual and RMSNorm fusions +// share this contraction without a separate dispatch or intermediate output. +func.def pure inline @ggml_q4_q8_lowtoken_token_pairs(%capacity: index, %tokens: index, %k: index, %n: index) -> (i1) { + %c4 = index.constant 4 : index + %c4096 = index.constant 4096 : index + %capacity4 = index.cmp eq, %capacity, %c4 : index + %tokens4 = index.cmp eq, %tokens, %c4 : index + %four_tokens = scalar.andi %capacity4, %tokens4 : i1 + %long_input = index.cmp uge, %k, %c4096 : index + %wide_output = index.cmp uge, %n, %c4096 : index + %contracts = index.cmp ule, %n, %k : index + %wide = scalar.andi %long_input, %wide_output : i1 + %contracting = scalar.andi %wide, %contracts : i1 + %eligible = scalar.andi %four_tokens, %contracting : i1 + func.return %eligible : i1 +} + +func.def inline @ggml_q4_q8_lowtoken_adjacent_scale_min(%scale0: i32, %scale1: i32, %scale2: i32, %pair0: index) -> (i32, i32, i32, i32) { + %pair = index.assume %pair0 [range(%pair0, 0, 3)] : index + %two = index.constant 2 : index + %sixteen = index.constant 16 : index + %low = index.cmp ult, %pair, %two : index + %pair_lane = index.rem %pair, %two : index + %shift_index = index.mul %pair_lane, %sixteen : index + %shift = index.cast %shift_index : index to i32 + %two_i32 = scalar.constant 2 : i32 + %four_i32 = scalar.constant 4 : i32 + %eight_i32 = scalar.constant 8 : i32 + %low_mask = scalar.constant 3855 : i32 + %high_mask = scalar.constant 12336 : i32 + %byte_mask = scalar.constant 255 : i32 + %high_shift = scalar.addi %shift, %two_i32 : i32 + %min_shift = scalar.addi %shift, %four_i32 : i32 + %s_source = scf.select %low, %scale0, %scale2 : i32 + %m_source = scf.select %low, %scale1, %scale2 : i32 + %s_high_shift = scf.select %low, %shift, %high_shift : i32 + %m_low_shift = scf.select %low, %shift, %min_shift : i32 + %s_lo_raw = scalar.shrui %s_source, %shift : i32 + %s_lo = scalar.andi %s_lo_raw, %low_mask : i32 + %s_hi_raw = scalar.shrui %scale0, %s_high_shift : i32 + %s_hi = scalar.andi %s_hi_raw, %high_mask : i32 + %s = scalar.ori %s_lo, %s_hi : i32 + %m_lo_raw = scalar.shrui %m_source, %m_low_shift : i32 + %m_lo = scalar.andi %m_lo_raw, %low_mask : i32 + %m_hi_raw = scalar.shrui %scale1, %s_high_shift : i32 + %m_hi = scalar.andi %m_hi_raw, %high_mask : i32 + %m = scalar.ori %m_lo, %m_hi : i32 + %s0 = scalar.andi %s, %byte_mask : i32 + %m0 = scalar.andi %m, %byte_mask : i32 + %s1 = scalar.shrui %s, %eight_i32 : i32 + %m1 = scalar.shrui %m, %eight_i32 : i32 + func.return %s0, %m0, %s1, %m1 : i32, i32, i32, i32 +} + +func.def inline @ggml_q4_q8_lowtoken_dot8(%paired_tokens: i1, %weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) { + %value = scf.if %paired_tokens -> (f32) { + %zero = vector.constant 0 : vector<1xi32> + %wi0 = vector.extract %weight[0] : vector<2xi32> -> i32 + %ai0 = vector.extract %activation[0] : vector<2xi32> -> i32 + %ww0 = vector.splat %wi0 : vector<1xi32> + %aw0 = vector.splat %ai0 : vector<1xi32> + %w0 = vector.bitcast %ww0 : vector<1xi32> to vector<4xi8> + %a0 = vector.bitcast %aw0 : vector<1xi32> to vector<4xi8> + %d0 = vector.dot4i %w0, %a0, %zero : vector<4xi8>, vector<4xi8>, vector<1xi32> + %wi1 = vector.extract %weight[1] : vector<2xi32> -> i32 + %ai1 = vector.extract %activation[1] : vector<2xi32> -> i32 + %ww1 = vector.splat %wi1 : vector<1xi32> + %aw1 = vector.splat %ai1 : vector<1xi32> + %w1 = vector.bitcast %ww1 : vector<1xi32> to vector<4xi8> + %a1 = vector.bitcast %aw1 : vector<1xi32> to vector<4xi8> + %d1 = vector.dot4i %w1, %a1, %d0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %sum = vector.extract %d1[0] : vector<1xi32> -> i32 + %result = scalar.sitofp %sum : i32 to f32 + scf.yield %result : f32 + } else { + %original = func.call @ggml_dot_u8_s8_vector8_f32(%weight, %activation) : (vector<2xi32>, vector<2xi32>) -> (f32) + scf.yield %original : f32 + } + func.return %value : f32 +} + +func.def inline @ggml_mul_mat_q4_q8_1_x4_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %paired_tokens = func.call pure @ggml_q4_q8_lowtoken_token_pairs(%capacity, %tokens, %k, %n) : (index, index, index, index) -> (i1) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %wg = index.add %channel_tile, %c0 : index + %c64_rows = index.constant 64 : index + %direct_rows = index.div %c8, %k_partitions : index + %row_tile = scf.select %staged, %c64_rows, %direct_rows : index + %row_base = index.mul %wg, %row_tile : index + %wave_rows = index.div %c4, %k_partitions : index + %wave_row = index.mul %wave, %wave_rows : index + %cohort_width = index.mul %c16, %k_partitions : index + %row_half = index.div %lane, %cohort_width : index + %cohort_lane = index.rem %lane, %cohort_width : index + %local_row = index.add %wave_row, %row_half : index + %row = index.add %row_base, %local_row : index + %valid_row = index.cmp ult, %row, %n : index + %cohort = index.div %lane, %c16 : index + %block_lane = index.rem %lane, %c16 : index + %pair = index.div %block_lane, %c4 : index + %packet = index.rem %block_lane, %c4 : index + %group0 = index.mul %pair, %c2 : index + %group1 = index.add %group0, %c1 : index + %word_page = index.mul %pair, %c8 : index + %word_add = index.mul %packet, %c2 : index + %word_index = index.add %word_page, %word_add : index + %meta_lane = index.cmp eq, %packet, %c0 : index + %blocks = index.div %k, %c256 : index + %partition_blocks = index.div %blocks, %k_partitions : index + %k_team = index.div %cohort_lane, %c16 : index + %k_begin = index.mul %k_team, %partition_blocks : index + %weight_row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %weight_row_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_groups = index.div %k, %c128 : index + %global_q8_row_bytes = index.scale %q8_groups, %block_bytes : index, offset -> offset + %slab_row_bytes = index.constant 2304 : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %q8_half = index.div %pair, %c2 : index + %q8_pair = index.rem %pair, %c2 : index + %q8_group0 = index.mul %q8_pair, %c2 : index + %q8_group1 = index.add %q8_group0, %c1 : index + %q8_words_base = index.mul %q8_group0, %c8 : index + %q8_word0 = index.add %q8_words_base, %word_add : index + %q8_word1 = index.add %q8_word0, %c8 : index + %q8_meta0 = index.mul %q8_group0, %c2 : index + %q8_meta1 = index.mul %q8_group1, %c2 : index + %q8_sum0 = index.add %q8_meta0, %c1 : index + %q8_sum1 = index.add %q8_meta1, %c1 : index + %c64 = index.constant 64 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %channel_tile, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %payload_field = index.div %word_index, %c4 : index + %payload_word = index.rem %word_index, %c4 : index + %input_na, %gate_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = index.constant 8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_header = vector.constant 0 : vector<4xi32> + %zero_raw = vector.constant 0 : vector<2xi32> + %prime_header, %prime_raw = scf.if %staged -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %c0, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %zero_header, %zero_raw : vector<4xi32>, vector<2xi32> + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_header, %last_raw = scf.for %block_base = [%c0 to %partition_blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_header = %prime_header : vector<4xi32>, %carry_raw = %prime_raw : vector<2xi32>) -> (f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32>) { + %block_in_slab = index.rem %block_base, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %new_slab = scalar.andi %staged, %slab_boundary : i1 + scf.if %new_slab { + func.call @ggml_q8_1_x4_stage_lowtoken_slab(%input_na, %q8_stage, %capacity, %block_base, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %block = index.add %block_base, %k_begin : index + %valid_block = index.cmp ult, %block, %blocks : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %group_block = index.add %row_group_block, %safe_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %g_header_view = buffer.view %gate_na[%wb] : buffer -> view<4xi32> + %g_code_view = buffer.view %gate_na[%wc] : buffer -> view<8x64x4xi32> + %g_headers = scf.if %staged -> (vector<4xi32>) { + scf.yield %carry_header : vector<4xi32> + } else { + %header = vector.load %g_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %header : vector<4xi32> + } + %g_dm_i32 = vector.extract %g_headers[0] : vector<4xi32> -> i32 + %g_dm_word = vector.splat %g_dm_i32 : vector<1xi32> + %g_dm = vector.bitcast %g_dm_word : vector<1xi32> to vector<2xf16> + %g_d0h = vector.extract %g_dm[0] : vector<2xf16> -> f16 + %g_d1h = vector.extract %g_dm[1] : vector<2xf16> -> f16 + %g_d0 = scalar.extf %g_d0h : f16 to f32 + %g_d1 = scalar.extf %g_d1h : f16 to f32 + %g_s0 = vector.extract %g_headers[1] : vector<4xi32> -> i32 + %g_s1 = vector.extract %g_headers[2] : vector<4xi32> -> i32 + %g_s2 = vector.extract %g_headers[3] : vector<4xi32> -> i32 + %pair_scale0, %pair_min0, %pair_scale1, %pair_min1 = scf.if %paired_tokens -> (i32, i32, i32, i32) { + %s0, %m0, %s1, %m1 = func.call @ggml_q4_q8_lowtoken_adjacent_scale_min(%g_s0, %g_s1, %g_s2, %pair) : (i32, i32, i32, index) -> (i32, i32, i32, i32) + scf.yield %s0, %m0, %s1, %m1 : i32, i32, i32, i32 + } else { + %zero_i32 = scalar.constant 0 : i32 + scf.yield %zero_i32, %zero_i32, %zero_i32, %zero_i32 : i32, i32, i32, i32 + } + %g_scale0i, %g_min0i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale0, %pair_min0 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale0f = scalar.uitofp %g_scale0i : i32 to f32 + %g_min0f = scalar.uitofp %g_min0i : i32 to f32 + %g_ws0 = scalar.mulf %g_d0, %g_scale0f : f32 + %g_wm0 = scalar.mulf %g_d1, %g_min0f : f32 + %g_scale1i, %g_min1i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale1, %pair_min1 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale1f = scalar.uitofp %g_scale1i : i32 to f32 + %g_min1f = scalar.uitofp %g_min1i : i32 to f32 + %g_ws1 = scalar.mulf %g_d0, %g_scale1f : f32 + %g_wm1 = scalar.mulf %g_d1, %g_min1f : f32 + %g_raw = scf.if %staged -> (vector<2xi32>) { + scf.yield %carry_raw : vector<2xi32> + } else { + %raw = vector.load %g_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %raw : vector<2xi32> + } + %g_mask = vector.constant 252645135 : vector<2xi32> + %g_shift = vector.constant 4 : vector<2xi32> + %g_lo = vector.andi %g_raw, %g_mask : vector<2xi32> + %g_high = vector.shrui %g_raw, %g_shift : vector<2xi32> + %g_hi = vector.andi %g_high, %g_mask : vector<2xi32> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_header, %next_raw = scf.if %has_next -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %next_block, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %carry_header, %carry_raw : vector<4xi32>, vector<2xi32> + } + %input_block = scf.select %staged, %block_in_slab, %safe_block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %q8_half : index + %q8_block_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next_g0 = scf.if %valid0 -> (f32) { + %tb = index.scale %t0, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ad0h = view.load %ds_view[%q8_meta0] : view<8xf16> -> f16 + %ad1h = view.load %ds_view[%q8_meta1] : view<8xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = view.load %ds_view[%q8_sum0] : view<8xf16> -> f16 + %as1h = view.load %ds_view[%q8_sum1] : view<8xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g0, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next_g1 = scf.if %valid1 -> (f32) { + %tb = index.scale %t1, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g1, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + scf.if %paired_tokens { + scf.schedule.fence + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next_g2 = scf.if %valid2 -> (f32) { + %tb = index.scale %t2, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g2, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next_g3 = scf.if %valid3 -> (f32) { + %tb = index.scale %t3, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g3, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next_g4 = scf.if %valid4 -> (f32) { + %tb = index.scale %t4, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g4, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_header, %next_raw : f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32> + } + %combined0 = func.call @ggml_lowtoken_combine_k_partitions(%total_g0, %k_partitions) : (f32, index) -> (f32) + %combined1 = func.call @ggml_lowtoken_combine_k_partitions(%total_g1, %k_partitions) : (f32, index) -> (f32) + %combined2 = func.call @ggml_lowtoken_combine_k_partitions(%total_g2, %k_partitions) : (f32, index) -> (f32) + %combined3 = func.call @ggml_lowtoken_combine_k_partitions(%total_g3, %k_partitions) : (f32, index) -> (f32) + %combined4 = func.call @ggml_lowtoken_combine_k_partitions(%total_g4, %k_partitions) : (f32, index) -> (f32) + func.call @ggml_mul_mat_q8_lowtoken_publish(%tokens, %n, %row_tile, %row_base, %local_row, %cohort_lane, %channel_tile, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %combined0, %combined1, %combined2, %combined3, %combined4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, index, index, index, f32, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Native Q6_K carrier with signed I8 dot products and F32 scaling. +func.def inline @ggml_q6_prefetch_native_words(%weight_na: buffer, %row_group_block: index, %block: index, %row_lane: index, %half16: index, %group: index, %half128: index, %ql_side: index, %packet: index) -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %c2 = index.constant 2 : index + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %group_block = index.add %row_group_block, %block : index + %block_base = index.scale %group_block, %record_bytes : index, offset -> offset + %scale_base = index.add %block_base, %scale_add : offset + %raw_base = index.add %block_base, %raw_add : offset + %d_view = buffer.view %weight_na[%block_base] : buffer -> view<64xf16> + %s_view = buffer.view %weight_na[%scale_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_na[%raw_base] : buffer -> view<2x3x64x8xi32> + %dh = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %si = view.load %s_view[%row_lane, %half16, %group] : view<64x2x8xi8> -> i8 + %ql = vector.load %raw_view[%half128, %ql_side, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + %qh = vector.load %raw_view[%half128, %c2, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + func.return %dh, %si, %ql, %qh : f16, i8, vector<4xi32>, vector<4xi32> +} + +func.def inline @ggml_q6_q8_lowtoken_activation_values(%pipelined: i1, %token: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index) -> (vector<16xi8>, f32) { + %code_offset = index.constant 16 : offset + %a, %ad = scf.if %pipelined -> (vector<16xi8>, f32) { + %base = index.constant 0 : offset + %c4 = index.constant 4 : index + %q8_block = index.assume %q8_block0 [range(%q8_block0, 0, 7)] : index + %word = index.add %q8_word0, %c4 : index + %q_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x36xi32> + %ds_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x72xf16> + %words = vector.load %q_view[%current_slot, %token, %q8_block, %word] : view<2x[%capacity]x8x36xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%current_slot, %token, %q8_block, %q8_meta] : view<2x[%capacity]x8x72xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } else { + %tb = index.scale %token, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %words = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%q8_meta] : view<8xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } + func.return %a, %ad : vector<16xi8>, f32 +} + +func.def inline @ggml_q6_q8_lowtoken_dot16(%pipelined: i1, %w: vector<16xi8>, %a: vector<16xi8>) -> (i32) { + %sum = scf.if %pipelined -> (i32) { + %di0 = vector.constant 0 : vector<1xi32> + %dw0 = vector.slice %w[0] : vector<16xi8> -> vector<4xi8> + %da0 = vector.slice %a[0] : vector<16xi8> -> vector<4xi8> + %di1 = vector.dot4i %dw0, %da0, %di0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw1 = vector.slice %w[4] : vector<16xi8> -> vector<4xi8> + %da1 = vector.slice %a[4] : vector<16xi8> -> vector<4xi8> + %di2 = vector.dot4i %dw1, %da1, %di1 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw2 = vector.slice %w[8] : vector<16xi8> -> vector<4xi8> + %da2 = vector.slice %a[8] : vector<16xi8> -> vector<4xi8> + %di3 = vector.dot4i %dw2, %da2, %di2 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw3 = vector.slice %w[12] : vector<16xi8> -> vector<4xi8> + %da3 = vector.slice %a[12] : vector<16xi8> -> vector<4xi8> + %di4 = vector.dot4i %dw3, %da3, %di3 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %chain_sum = vector.extract %di4[0] : vector<1xi32> -> i32 + scf.yield %chain_sum : i32 + } else { + %init = vector.constant 0 : vector<4xi32> + %dot = vector.dot4i %w, %a, %init : vector<16xi8>, vector<16xi8>, vector<4xi32> + %seed = scalar.constant 0 : i32 + %direct_sum = vector.reduce %dot, %seed : vector<4xi32>, i32 + scf.yield %direct_sum : i32 + } + func.return %sum : i32 +} + +func.def inline @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input: buffer, %capacity0: index, %slot_begin: index, %blocks: index, %global_row_bytes: offset) -> (vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c18 = index.constant 18 : index + %c72 = index.constant 72 : index + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %copy_packet = index.rem %thread, %total_packets : index + %copy_token = index.div %copy_packet, %c72 : index + %copy_inner = index.rem %copy_packet, %c72 : index + %remaining = index.sub %blocks, %slot_begin : index + %full = index.cmp uge, %remaining, %c4 : index + %active0 = scf.select %full, %c4, %remaining : index + %active = index.assume %active0 [range(%active0, 1, 4)] : index + %active_packets = index.mul %active, %c18 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + %safe_inner = scf.select %copy_valid, %copy_inner, %c0 : index + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %slot_bytes = index.scale %slot_begin, %block_bytes : index, offset -> offset + %copy_bytes = index.scale %safe_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slot_bytes : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + func.return %values : vector<4xi32> +} +func.def inline @ggml_q8_1_x4_publish_lowtoken_half_slab(%stage: buffer, %capacity0: index, %slot0: index, %values: vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %slot = index.assume %slot0 [range(%slot0, 0, 1)] : index + %c4 = index.constant 4 : index + %c72 = index.constant 72 : index + %base = index.constant 0 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %thread_valid = index.cmp ult, %thread, %total_packets : index + scf.if %thread_valid { + %copy_token0 = index.div %thread, %c72 : index + %copy_token = index.assume %copy_token0 [range(%copy_token0, 0, 4), lt(%copy_token0, %capacity)] : index + %copy_inner = index.rem %thread, %c72 : index + %origin = index.mul %copy_inner, %c4 : index + %destination = buffer.view %stage[%base] : buffer -> view<2x[%capacity]x288xi32> + vector.store %values, %destination[%slot, %copy_token, %origin] : vector<4xi32>, view<2x[%capacity]x288xi32> + } + func.return +} + + +func.def inline @ggml_q6_q8_lowtoken_accumulate(%pipelined: i1, %tokens: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index, %w: vector<16xi8>, %ws: f32, %acc_g0: f32, %acc_g1: f32, %acc_g2: f32, %acc_g3: f32, %acc_g4: f32) -> (f32, f32, f32, f32, f32) { + %code_offset = index.constant 16 : offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %computed_g0 = scf.if %valid0 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t0, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g0 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %computed_g1 = scf.if %valid1 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t1, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g1 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %computed_g2 = scf.if %valid2 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t2, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g2 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %computed_g3 = scf.if %valid3 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t3, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g3 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %computed_g4 = scf.if %valid4 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t4, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g4 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + func.return %computed_g0, %computed_g1, %computed_g2, %computed_g3, %computed_g4 : f32, f32, f32, f32, f32 +} + +func.def inline @ggml_mul_mat_q6_q8_1_x4_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %row_tile = scf.select %staged, %c32, %c8 : index + %row_base = index.mul %channel_tile, %row_tile : index + %wave_row = index.mul %wave, %c4 : index + %cohort = index.div %lane, %c16 : index + %local_row = index.add %wave_row, %cohort : index + %row = index.add %row_base, %local_row : index + %block_lane = index.rem %lane, %c16 : index + %group = index.div %block_lane, %c2 : index + %half16 = index.rem %block_lane, %c2 : index + %half128 = index.div %group, %c4 : index + %group_in_half = index.rem %group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift0 = index.mul %nibble, %c4 : index + %nibble_shift1 = index.cast %nibble_shift0 : index to i32 + %nibble_shift = vector.splat %nibble_shift1 : vector<4xi32> + %high_shift0 = index.mul %group_in_half, %c2 : index + %high_shift1 = index.cast %high_shift0 : index to i32 + %high_shift = vector.splat %high_shift1 : vector<4xi32> + %packet = index.mul %half16, %c4 : index + %q8_group_words = index.mul %group_in_half, %c8 : index + %q8_word0 = index.add %q8_group_words, %packet : index + %q8_meta = index.mul %group_in_half, %c2 : index + %blocks = index.div %k, %c256 : index + %q8_groups = index.div %k, %c128 : index + %q8_block_bytes = index.constant 144 : offset + %global_q8_row_bytes = index.scale %q8_groups, %q8_block_bytes : index, offset -> offset + %native_block_bytes = index.constant 210 : index + %cache_bytes = index.constant 33554432 : index + %weight_blocks = index.mul %blocks, %n : index + %weight_bytes = index.mul %weight_blocks, %native_block_bytes : index + %exceeds_cache = index.cmp uge, %weight_bytes, %cache_bytes : index + %pipelined = scalar.andi %staged, %exceeds_cache : i1 + %full_slab_row_bytes = index.constant 2304 : offset + %half_slab_row_bytes = index.constant 1152 : offset + %slab_row_bytes = scf.select %pipelined, %half_slab_row_bytes, %full_slab_row_bytes : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %code_offset = index.constant 16 : offset + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %staged_row_group = index.div %channel_tile, %c2 : index + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %staged_row_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %input_na, %weight_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = scf.select %pipelined, %c4, %c8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_half = scalar.constant 0.0 : f16 + %zero_code = scalar.constant 0 : i8 + %zero_words = vector.constant 0 : vector<4xi32> + %prime_d, %prime_s, %prime_ql, %prime_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %pd, %ps, %pql, %pqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %c0, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %pd, %ps, %pql, %pqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %zero_half, %zero_code, %zero_words, %zero_words : f16, i8, vector<4xi32>, vector<4xi32> + } + scf.if %pipelined { + %prime_q8 = func.call @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input_na, %capacity, %c0, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + func.call @ggml_q8_1_x4_publish_lowtoken_half_slab(%q8_stage, %capacity, %c0, %prime_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_d, %last_s, %last_ql, %last_qh = scf.for %block = [%c0 to %blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_d = %prime_d : f16, %carry_s = %prime_s : i8, %carry_ql = %prime_ql : vector<4xi32>, %carry_qh = %prime_qh : vector<4xi32>) -> (f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32>) { + %block_in_slab = index.rem %block, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %direct_stage = scalar.andi %staged, %slab_boundary : i1 + %true = scalar.constant true : i1 + %not_pipelined = scalar.xori %pipelined, %true : i1 + %stage_current = scalar.andi %direct_stage, %not_pipelined : i1 + scf.if %stage_current { + func.call @ggml_q8_1_x4_stage_lowtoken_slab(%input_na, %q8_stage, %capacity, %block, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %slab_number = index.div %block, %slab_blocks : index + %current_slot = index.rem %slab_number, %c2 : index + %following_slab = index.add %slab_number, %c1 : index + %following_begin = index.mul %following_slab, %slab_blocks : index + %following_exists = index.cmp ult, %following_begin, %blocks : index + %next_slot = index.rem %following_slab, %c2 : index + %staged_future = scalar.andi %pipelined, %following_exists : i1 + %current_d, %current_s, %current_ql, %current_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + %wd, %ws8, %wql, %wqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } + %d = scalar.extf %current_d : f16 to f32 + %sf = scalar.sitofp %current_s : i8 to f32 + %ws = scalar.mulf %d, %sf : f32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %four = vector.constant 4 : vector<4xi32> + %ql_shifted = vector.shrui %current_ql, %nibble_shift : vector<4xi32> + %low = vector.andi %ql_shifted, %nibble_mask : vector<4xi32> + %qh_shifted = vector.shrui %current_qh, %high_shift : vector<4xi32> + %high0 = vector.andi %qh_shifted, %high_mask : vector<4xi32> + %high = vector.shli %high0, %four : vector<4xi32> + %code = vector.ori %low, %high : vector<4xi32> + // Guard bits prevent subtraction from borrowing across packed bytes. + %guard_mask = vector.constant -2139062144 : vector<4xi32> + %center = vector.constant 538976288 : vector<4xi32> + %guarded = vector.ori %code, %guard_mask : vector<4xi32> + %centered = vector.subi %guarded, %center : vector<4xi32> + %signed_code = vector.xori %centered, %guard_mask : vector<4xi32> + %w = vector.bitcast %signed_code : vector<4xi32> to vector<16xi8> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_d, %next_s, %next_ql, %next_qh = scf.if %has_next -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %wd, %ws8, %wql, %wqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %next_block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } + %input_block = scf.select %staged, %block_in_slab, %block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %half128 : index + %q8_local_block_add = index.scale %q8_block, %q8_block_bytes : index, offset -> offset + %slot_row = index.mul %current_slot, %capacity : index + %slot_byte_add0 = index.scale %slot_row, %slab_row_bytes : index, offset -> offset + %slot_byte_add = scf.select %pipelined, %slot_byte_add0, %base : offset + %q8_block_add = index.add %slot_byte_add, %q8_local_block_add : offset + %slot_complete = index.cmp eq, %block_in_slab, %c3 : index + %publish_next = scalar.andi %staged_future, %slot_complete : i1 + %next_g0, %next_g1, %next_g2, %next_g3, %next_g4 = scf.if %publish_next -> (f32, f32, f32, f32, f32) { + %next_q8 = func.call @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input_na, %capacity, %following_begin, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_q6_q8_lowtoken_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + func.call @ggml_q8_1_x4_publish_lowtoken_half_slab(%q8_stage, %capacity, %next_slot, %next_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } else { + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_q6_q8_lowtoken_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_d, %next_s, %next_ql, %next_qh : f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32> + } + func.call @ggml_mul_mat_q8_lowtoken_publish(%tokens, %n, %row_tile, %row_base, %local_row, %block_lane, %channel_tile, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, index, index, index, f32, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %activation_format = config.get @ggml.mul_mat.activation_format : index + %q8_format = index.constant 9 : index + %uses_q8 = index.cmp eq, %activation_format, %q8_format : index + scf.if %uses_q8 { + %q6_format = index.constant 46 : index + %uses_q6 = index.cmp eq, %weight_format, %q6_format : index + scf.if %uses_q6 { + func.call @ggml_mul_mat_q6_q8_1_x4_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_mul_mat_q4_q8_1_x4_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } else { + func.call @ggml_mul_mat_f32_f32_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token2(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 2 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token3(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 3 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token4(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 4 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token5(%weight_format: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.constant 5 : index + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_lowtoken_dispatch(%weight_format, %tokens, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %tokens = index.assume %token_count [range(%token_count, 2, 5)] : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %is_token2 = index.cmp eq, %tokens, %c2 : index + scf.if %is_token2 { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token2(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + %is_token3 = index.cmp eq, %tokens, %c3 : index + scf.if %is_token3 { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token3(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + %is_token4 = index.cmp eq, %tokens, %c4 : index + scf.if %is_token4 { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token4(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_matmul_skinny_input_f32_dequant_accumulate_f32_token5(%weight_format, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epilogue, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom new file mode 100644 index 000000000000..24d2078df656 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_postprocess_f32.loom @@ -0,0 +1,31 @@ +template.decl @ggml.matmul.skinny_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.skinny_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.def<@ggml.matmul.skinny_postprocess_f32.add_bias_vector4> device @ggml_matmul_skinny_postprocess_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.skinny_postprocess_f32.add_residual_vector4> device @ggml_matmul_skinny_postprocess_f32_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom new file mode 100644 index 000000000000..519714efb716 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/skinny_publish_f32.loom @@ -0,0 +1,14 @@ +template.decl @ggml.matmul.skinny_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul.skinny_publish_f32.store_vector4> device @ggml_matmul_skinny_publish_f32_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/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 new file mode 100644 index 000000000000..a4b588739396 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,237 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.mul_mat.activation_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Shared contiguous 2D F32 activation tiled matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4 +// ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%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(%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, 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 + %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 + %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 + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %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<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 + %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 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %epilogue, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + template.apply<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile>(%bounded_token_count, %bounded_output_size, %output_tile_count, %token_tile_count, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_input_f32_dequant_accumulate_f32_core(%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 new file mode 100644 index 000000000000..42fc4e44b4f7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,237 @@ +// 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(%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) + +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%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 + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %f32_element_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %gate_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%gate_weight_format, %bounded_input_size) : (index, index) -> (offset) + %up_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%up_weight_format, %bounded_input_size) : (index, index) -> (offset) + %input_noalias, %gate_weight_noalias, %up_weight_noalias = buffer.assume.noalias %input, %gate_weight, %up_weight : buffer, buffer, buffer + %gate_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %up_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %gate_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %up_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %gate_weight_stage_view = buffer.view %gate_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %up_weight_stage_view = buffer.view %up_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %gate_result_fragment_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %up_result_fragment_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %gate_result_physical_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %up_result_physical_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %gate_result00, %gate_result01, %gate_result10, %gate_result11, %up_result00, %up_result01, %up_result10, %up_result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%gate_block_acc00 = %init00 : vector<8xf16>, %gate_block_acc01 = %init01 : vector<8xf16>, %gate_block_acc10 = %init10 : vector<8xf16>, %gate_block_acc11 = %init11 : vector<8xf16>, %up_block_acc00 = %init00 : vector<8xf16>, %up_block_acc01 = %init01 : vector<8xf16>, %up_block_acc10 = %init10 : vector<8xf16>, %up_block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%gate_acc00 = %gate_block_acc00 : vector<8xf16>, %gate_acc01 = %gate_block_acc01 : vector<8xf16>, %gate_acc10 = %gate_block_acc10 : vector<8xf16>, %gate_acc11 = %gate_block_acc11 : vector<8xf16>, %up_acc00 = %up_block_acc00 : vector<8xf16>, %up_acc01 = %up_block_acc01 : vector<8xf16>, %up_acc10 = %up_block_acc10 : vector<8xf16>, %up_acc11 = %up_block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %gate_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %gate_weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %up_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %up_weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %gate_weight_values, %gate_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + vector.store %up_weight_values, %up_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %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 + %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_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 { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 = scf.for %k_half = [%c0 to %c32 step %c16](%gate_half_acc00 = %gate_acc00 : vector<8xf16>, %gate_half_acc01 = %gate_acc01 : vector<8xf16>, %gate_half_acc10 = %gate_acc10 : vector<8xf16>, %gate_half_acc11 = %gate_acc11 : vector<8xf16>, %up_half_acc00 = %up_acc00 : vector<8xf16>, %up_half_acc01 = %up_acc01 : vector<8xf16>, %up_half_acc10 = %up_acc10 : vector<8xf16>, %up_half_acc11 = %up_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %gate_lhs0 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %gate_lhs1 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs0 = vector.fragment.load %up_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs1 = vector.fragment.load %up_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %gate_half_next00 = vector.mma %gate_lhs0, %rhs0, %gate_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next01 = vector.mma %gate_lhs0, %rhs1, %gate_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next10 = vector.mma %gate_lhs1, %rhs0, %gate_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next11 = vector.mma %gate_lhs1, %rhs1, %gate_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next00 = vector.mma %up_lhs0, %rhs0, %up_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next01 = vector.mma %up_lhs0, %rhs1, %up_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next10 = vector.mma %up_lhs1, %rhs0, %up_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next11 = vector.mma %up_lhs1, %rhs1, %up_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %gate_half_next00, %gate_half_next01, %gate_half_next10, %gate_half_next11, %up_half_next00, %up_half_next01, %up_half_next10, %up_half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %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 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + vector.fragment.store %gate_result00, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result00, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result01, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result01, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result10, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result10, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result11, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result11, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %op, %epilogue, %gate_wide, %up_wide, %bias, %output, %residual_input, %residual_output) : (i1, index, index, index, index, index, index, vector<4xf32>, vector<4xf32>, buffer, buffer, buffer, buffer) + } + func.return +} + +template.def<@ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core> device @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32_core(%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_pair_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_postprocess_f32.loom new file mode 100644 index 000000000000..ed61ece2dc27 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_postprocess_f32.loom @@ -0,0 +1,84 @@ +// Pairwise tiled matmul postprocess for gate/up projections. +template.decl @ggml.matmul.tiled_pair_postprocess_f32.apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +template.def<@ggml.matmul.tiled_pair_postprocess_f32.apply_vector4> device @ggml_matmul_tiled_pair_postprocess_f32_apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) { + %c0_5_f32x4 = vector.constant 0.5 : vector<4xf32> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c1_f32x4 = vector.constant 1.0 : vector<4xf32> + %c2_f32x4 = vector.constant 2.0 : vector<4xf32> + %cn2_f32x4 = vector.constant -2.0 : vector<4xf32> + %gelu_coef = vector.constant 0.044715 : vector<4xf32> + %sqrt_2_over_pi = vector.constant 0.79788456080286544 : vector<4xf32> + + %op_add = index.constant 0 : index + %sum = vector.addf %lhs, %rhs : vector<4xf32> + %is_add = index.cmp eq, %op, %op_add : index + %add_selected = scf.select %is_add, %sum, %lhs : vector<4xf32> + + %op_sub = index.constant 1 : index + %difference = vector.subf %lhs, %rhs : vector<4xf32> + %is_sub = index.cmp eq, %op, %op_sub : index + %sub_selected = scf.select %is_sub, %difference, %add_selected : vector<4xf32> + + %op_mul = index.constant 2 : index + %product = vector.mulf %lhs, %rhs : vector<4xf32> + %is_mul = index.cmp eq, %op, %op_mul : index + %mul_selected = scf.select %is_mul, %product, %sub_selected : vector<4xf32> + + %op_div = index.constant 3 : index + %quotient = vector.divf %lhs, %rhs : vector<4xf32> + %is_div = index.cmp eq, %op, %op_div : index + %div_selected = scf.select %is_div, %quotient, %mul_selected : vector<4xf32> + + %op_swiglu = index.constant 4 : index + %negative_lhs = vector.subf %c0_f32x4, %lhs : vector<4xf32> + %lhs_exp = vector.expf %negative_lhs : vector<4xf32> + %denominator = vector.addf %c1_f32x4, %lhs_exp : vector<4xf32> + %sigmoid = vector.divf %c1_f32x4, %denominator : vector<4xf32> + %silu = vector.mulf %lhs, %sigmoid : vector<4xf32> + %swiglu = vector.mulf %silu, %rhs : vector<4xf32> + %is_swiglu = index.cmp eq, %op, %op_swiglu : index + %swiglu_selected = scf.select %is_swiglu, %swiglu, %div_selected : vector<4xf32> + + %op_geglu = index.constant 5 : index + %gelu_x2 = vector.mulf %lhs, %lhs : vector<4xf32> + %gelu_poly0 = vector.mulf %gelu_coef, %gelu_x2 : vector<4xf32> + %gelu_poly = vector.addf %c1_f32x4, %gelu_poly0 : vector<4xf32> + %gelu_inner0 = vector.mulf %lhs, %gelu_poly : vector<4xf32> + %gelu_inner = vector.mulf %sqrt_2_over_pi, %gelu_inner0 : vector<4xf32> + %gelu_tanh_scaled = vector.mulf %cn2_f32x4, %gelu_inner : vector<4xf32> + %gelu_tanh_exp = vector.expf %gelu_tanh_scaled : vector<4xf32> + %gelu_tanh_denominator = vector.addf %c1_f32x4, %gelu_tanh_exp : vector<4xf32> + %gelu_tanh_ratio = vector.divf %c2_f32x4, %gelu_tanh_denominator : vector<4xf32> + %gelu_tanh = vector.subf %gelu_tanh_ratio, %c1_f32x4 : vector<4xf32> + %gelu_one_plus = vector.addf %c1_f32x4, %gelu_tanh : vector<4xf32> + %gelu_half_x = vector.mulf %c0_5_f32x4, %lhs : vector<4xf32> + %gelu = vector.mulf %gelu_half_x, %gelu_one_plus : vector<4xf32> + %geglu = vector.mulf %gelu, %rhs : vector<4xf32> + %is_geglu = index.cmp eq, %op, %op_geglu : index + %geglu_selected = scf.select %is_geglu, %geglu, %swiglu_selected : vector<4xf32> + + %op_reglu = index.constant 6 : index + %unary_relu = index.constant 2 : index + %relu_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_relu, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %reglu = vector.mulf %relu_lhs, %rhs : vector<4xf32> + %is_reglu = index.cmp eq, %op, %op_reglu : index + %reglu_selected = scf.select %is_reglu, %reglu, %geglu_selected : vector<4xf32> + + %op_geglu_erf = index.constant 7 : index + %unary_gelu_erf = index.constant 20 : index + %gelu_erf_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_gelu_erf, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %geglu_erf = vector.mulf %gelu_erf_lhs, %rhs : vector<4xf32> + %is_geglu_erf = index.cmp eq, %op, %op_geglu_erf : index + %geglu_erf_selected = scf.select %is_geglu_erf, %geglu_erf, %reglu_selected : vector<4xf32> + + %op_geglu_quick = index.constant 8 : index + %unary_gelu_quick = index.constant 14 : index + %gelu_quick_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_gelu_quick, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %geglu_quick = vector.mulf %gelu_quick_lhs, %rhs : vector<4xf32> + %is_geglu_quick = index.cmp eq, %op, %op_geglu_quick : index + %result = scf.select %is_geglu_quick, %geglu_quick, %geglu_erf_selected : vector<4xf32> + template.return %result : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom new file mode 100644 index 000000000000..bd5ec88c580a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_postprocess_f32.loom @@ -0,0 +1,31 @@ +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.def<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4> device @ggml_matmul_tiled_postprocess_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4> device @ggml_matmul_tiled_postprocess_f32_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %f32_element_bytes = index.constant 4 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %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_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 new file mode 100644 index 000000000000..b03b98130939 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_publish.loom @@ -0,0 +1,44 @@ +template.decl @ggml.matmul.tiled_publish.store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul.tiled_publish.store_vector4> device @ggml_matmul_tiled_publish_store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + %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 + %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 + %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 { + %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 { + %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_row[%channel], %mask : vector<4xf16>, view<[%output_size]xf16>, vector<4xi1> + } + template.return +} + +template.def<@ggml.matmul.tiled_publish_f32.store_vector4> device @ggml_matmul_tiled_publish_f32_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %f32_format = index.constant 32 : index + %publish_word = scalar.constant true : i1 + %c0_offset = index.constant 0 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%f32_format, %publish_word, %token_count0, %output_size0, %token0, %channel, %values, %scratch_values, %scratch_d, %output) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom new file mode 100644 index 000000000000..60e0a5b34e53 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_activation.loom @@ -0,0 +1,1470 @@ +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +config.def @ggml.matmul.vector.activation_format = 0 : index +config.def @ggml.matmul.vector.publish_format = 0 : index + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +// Shared vector-shaped matmul reduction pipeline. +// +// Op wrappers provide publication and output-group hooks by defining: +// ggml.matmul.vector_activation.publish_vector4 +// ggml.matmul.vector_activation.finish_output_group +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: i1, %arg6: f32, %arg7: f32, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +func.decl @ggml_matmul_vector_f32_decode_pair(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + + +func.def pure inline @ggml_matmul_vector_q8_geometry(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %q4_format = index.constant 44 : index + %q6_format = index.constant 46 : index + %q8_format = index.constant 9 : index + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %native_q4 = index.cmp eq, %weight_format, %q4_format : index + %native_q6 = index.cmp eq, %weight_format, %q6_format : index + %q8_input = index.cmp eq, %activation_format, %q8_format : index + %q4_min_blocks = index.constant 327680 : index + %blocks = index.div %input_size, %c256 : index + %weight_blocks = index.mul %blocks, %output_size : index + %large_payload = index.cmp uge, %weight_blocks, %q4_min_blocks : index + %q4_can_stage = scalar.andi %native_q4, %large_payload : i1 + %native_staged = scalar.ori %q4_can_stage, %native_q6 : i1 + %native_q8 = scalar.andi %native_staged, %q8_input : i1 + %q4_min_k = index.constant 16384 : index + %q6_min_k = index.constant 8192 : index + %min_k = scf.select %native_q4, %q4_min_k, %q6_min_k : index + %min_n = index.constant 4096 : index + %five_tokens = index.cmp eq, %token_capacity, %c5 : index + %long_k = index.cmp uge, %input_size, %min_k : index + %wide_n = index.cmp uge, %output_size, %min_n : index + %contracts = index.cmp ule, %output_size, %input_size : index + %large = scalar.andi %long_k, %wide_n : i1 + %long_reduction = scalar.andi %large, %contracts : i1 + %staged_shape = scalar.andi %five_tokens, %long_reduction : i1 + %staged = scalar.andi %native_q8, %staged_shape : i1 + %rows = scf.select %native_q4, %c64, %c32 : index + %result = scf.select %staged, %rows, %c0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %block_bytes = index.constant 144 : index + %half_cache_bytes = index.constant 16777216 : index + %cache_bytes = index.constant 33554432 : index + %payload_bytes = index.mul %weight_blocks, %block_bytes : index + %large_half = index.cmp uge, %payload_bytes, %half_cache_bytes : index + %fits_cache = index.cmp ule, %payload_bytes, %cache_bytes : index + %cache_range = scalar.andi %large_half, %fits_cache : i1 + %remainder = index.rem %blocks, %c2 : index + %even_blocks = index.cmp eq, %remainder, %c0 : index + %reduces_rows = index.cmp ult, %output_size, %input_size : index + %native_q4_q8 = scalar.andi %native_q4, %q8_input : i1 + %plain = index.cmp eq, %result, %c0 : index + %q4_plain = scalar.andi %native_q4_q8, %plain : i1 + %five_wide = scalar.andi %five_tokens, %wide_n : i1 + %reduction = scalar.andi %reduces_rows, %even_blocks : i1 + %partition_shape = scalar.andi %five_wide, %reduction : i1 + %partition_payload = scalar.andi %partition_shape, %cache_range : i1 + %partitioned = scalar.andi %q4_plain, %partition_payload : i1 + %k_partitions = scf.select %partitioned, %c2, %c1 : index + func.return %result, %k_partitions : index, index +} + +func.def pure inline @ggml_matmul_vector_uses_dot(%token_capacity: index, %input_size: index) -> (i1) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %few_tokens = index.cmp ule, %token_capacity, %c5 : index + %tail = index.rem %input_size, %c32 : index + %aligned_reduction = index.cmp eq, %tail, %c0 : index + %use_dot = scalar.andi %few_tokens, %aligned_reduction : i1 + func.return %use_dot : i1 +} + +func.def pure inline @ggml_matmul_vector_uses_f32_decode_pair(%token_capacity: index, %input_size: index, %publish_format: index) -> (i1) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %one_token = index.cmp eq, %token_capacity, %c1 : index + %tail = index.rem %input_size, %c32 : index + %aligned_reduction = index.cmp eq, %tail, %c0 : index + %plain_publish = index.cmp eq, %publish_format, %c0 : index + %shape = scalar.andi %one_token, %aligned_reduction : i1 + %f32_shape = scalar.andi %shape, %f32_input : i1 + %use_pair = scalar.andi %f32_shape, %plain_publish : i1 + func.return %use_pair : i1 +} + +template.def<@ggml.matmul.vector_activation.launch> pure @ggml_matmul_vector_activation_launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_groups = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_groups = index.div %padded_token_count, %c32 : index + %c5 = index.constant 5 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %vector_reduction = func.call pure @ggml_matmul_vector_uses_dot(%token_capacity, %input_size) : (index, index) -> (i1) + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %small_output = index.cmp ult, %output_size, %c64 : index + %narrow = scalar.andi %vector_reduction, %small_output : i1 + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %staged_rows, %k_partitions = func.call pure @ggml_matmul_vector_q8_geometry(%token_capacity, %input_size, %output_size, %weight_format) : (index, index, index, index) -> (index, index) + %staged = index.cmp ugt, %staged_rows, %c0 : index + %staged_rounding = index.sub %staged_rows, %c1 : index + %staged_threads = index.mul %staged_rows, %c16 : index + %cohort_rows = index.div %c8, %k_partitions : index + %cohort_rounding = index.sub %cohort_rows, %c1 : index + %base_rows_per_group = scf.select %narrow, %c4, %cohort_rows : index + %c1024 = index.constant 1024 : index + %c4096 = index.constant 4096 : index + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %batched_tokens = index.cmp ugt, %token_capacity, %c1 : index + %long_reduction = index.cmp uge, %input_size, %c4096 : index + %partition_tail = index.rem %input_size, %c1024 : index + %full_partitions = index.cmp eq, %partition_tail, %c0 : index + %batch_f32 = scalar.andi %batched_tokens, %f32_input : i1 + %partitionable = scalar.andi %long_reduction, %full_partitions : i1 + %partition_rows = scalar.andi %batch_f32, %partitionable : i1 + %narrow_threads = scf.select %partition_rows, %c1024, %c256 : index + %base_workgroup_size = scf.select %narrow, %narrow_threads, %c128 : index + %base_rounding = scf.select %narrow, %c3, %cohort_rounding : index + %row_rounding = scf.select %staged, %staged_rounding, %base_rounding : index + %rows_per_group = scf.select %staged, %staged_rows, %base_rows_per_group : index + %workgroup_size = scf.select %staged, %staged_threads, %base_workgroup_size : index + %rounded_n = index.add %output_size, %row_rounding : index + %direct_row_groups = index.div %rounded_n, %rows_per_group : index + %publish_format = config.get @ggml.matmul.vector.publish_format : index + %pair_decode = func.call pure @ggml_matmul_vector_uses_f32_decode_pair(%token_capacity, %input_size, %publish_format) : (index, index, index) -> (i1) + %padded_pair_output = index.add %output_size, %c1 : index + %pair_groups = index.div %padded_pair_output, %c2 : index + %pair_workgroup_size = index.constant 64 : index + %reduction_launch_x = scf.select %pair_decode, %pair_groups, %direct_row_groups : index + %reduction_workgroup_size = scf.select %pair_decode, %pair_workgroup_size, %workgroup_size : index + %launch_x = scf.select %vector_reduction, %reduction_launch_x, %output_groups : index + %launch_y = scf.select %vector_reduction, %c1, %token_groups : index + template.return %launch_x, %launch_y, %c1, %reduction_workgroup_size : index, index, index, index +} + +func.def inline @ggml_matmul_vector_combine_k_partitions(%value: f32, %partitions: index) -> (f32) { + %c1 = index.constant 1 : index + %partitioned = index.cmp ugt, %partitions, %c1 : index + %result = scf.if %partitioned -> (f32) { + %offset = scalar.constant 16 : i32 + %width = scalar.constant 32 : i32 + %peer, %ok = kernel.subgroup.shuffle %value, %offset, %width : f32, i32, i32 + %sum = scalar.addf %value, %peer : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_matmul_vector_reduce_cohort_f32(%value: f32) -> (f32) { + %width = scalar.constant 32 : i32 + %off1 = scalar.constant 1 : i32 + %peer1, %ok1 = kernel.subgroup.shuffle %value, %off1, %width : f32, i32, i32 + %sum1 = scalar.addf %value, %peer1 : f32 + %off2 = scalar.constant 2 : i32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %off2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %off4 = scalar.constant 4 : i32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %off4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %off8 = scalar.constant 8 : i32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %off8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + func.return %sum8 : f32 +} + +// Narrow vector projections assign a full wave to each row. +// Both cohort sizes use the existing wrapper publication hooks. +func.def inline @ggml_matmul_vector_reduce_f32(%values: vector<4xf32>, %full_wave: i1) -> (f32) { + %zero = scalar.constant 0.0 : f32 + %width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %sum = vector.reduce %values, %zero : vector<4xf32>, f32 + %peer1, %ok1 = kernel.subgroup.shuffle %sum, %xor1, %width : f32, i32, i32 + %sum1 = scalar.addf %sum, %peer1 : f32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %xor2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %xor4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %xor8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + %result = scf.if %full_wave -> (f32) { + %xor16 = scalar.constant 16 : i32 + %peer16, %ok16 = kernel.subgroup.shuffle %sum8, %xor16, %width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer16 : f32 + %lane32 = scalar.constant 32 : i32 + %peer32 = kernel.subgroup.broadcast %sum16 from %lane32 : f32, i32 + %sum64 = scalar.addf %sum16, %peer32 : f32 + scf.yield %sum64 : f32 + } else { + scf.yield %sum8 : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_matmul_vector_f32_activation_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %tokens = index.assume %token_count [range(%token_count, 1, 1)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %workgroup_size = kernel.workgroup.size : index + %one_wave_per_row = index.cmp eq, %workgroup_size, %c256 : index + %c1024 = index.constant 1024 : index + %partition_group = index.cmp eq, %workgroup_size, %c1024 : index + %narrow_output = index.cmp ult, %n, %c64 : index + %partitioned = scalar.andi %partition_group, %narrow_output : i1 + %full_wave = scalar.ori %one_wave_per_row, %partitioned : i1 + %k_partitions = scf.select %partitioned, %c4, %c1 : index + %cohort_width = scf.select %full_wave, %c64, %c16 : index + %rows_per_group = scf.select %full_wave, %c4, %c8 : index + %wave_rows = scf.select %full_wave, %c1, %c4 : index + %step_width = scf.select %full_wave, %c256, %c64 : index + %steps_per_block = scf.select %full_wave, %c1, %c4 : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %lane16 = index.rem %lane, %cohort_width : index + %cohort = index.div %lane, %cohort_width : index + %row_group_base = index.mul %channel_group, %rows_per_group : index + %row_subgroup = index.div %subgroup, %k_partitions : index + %k_partition = index.rem %subgroup, %k_partitions : index + %wave_add = index.mul %row_subgroup, %wave_rows : index + %wave_base = index.add %row_group_base, %wave_add : index + %row = index.add %wave_base, %cohort : index + %valid_row = index.cmp ult, %row, %n : index + %packet = index.rem %lane16, %c8 : index + %lane_group = index.div %lane16, %c8 : index + %lane_k = index.mul %lane16, %c4 : index + %c255 = index.constant 255 : index + %padded_k = index.add %k, %c255 : index + %blocks = index.div %padded_k, %c256 : index + %step_rounding = index.sub %step_width, %c1 : index + %padded_steps = index.add %k, %step_rounding : index + %steps = index.div %padded_steps, %step_width : index + %partition_steps = index.div %steps, %k_partitions : index + %step_begin = index.mul %k_partition, %partition_steps : index + %step_end = index.add %step_begin, %partition_steps : index + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %k) : (index, index) -> (offset) + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %a_na, %w_na = buffer.assume.noalias %input, %weight : buffer, buffer + %av = buffer.view %a_na[%base] : buffer -> view<[%tokens]x[%blocks]x256xf32> + %total0, %total1, %total2, %total3, %total4 = scf.for %step = [%step_begin to %step_end step %c1](%acc0 = %zero : vector<4xf32>, %acc1 = %zero : vector<4xf32>, %acc2 = %zero : vector<4xf32>, %acc3 = %zero : vector<4xf32>, %acc4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block = index.div %step, %steps_per_block : index + %quarter = index.rem %step, %steps_per_block : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %step_k = index.mul %step, %step_width : index + %kk = index.add %step_k, %lane_k : index + %valid_k = index.cmp ult, %kk, %k : index + %valid_row_k = scalar.andi %valid_row, %valid_k : i1 + %local_step_k = index.mul %quarter, %step_width : index + %local_k = index.add %local_step_k, %lane_k : index + %input_mask = vector.mask.range [%kk to %k step %c1] : index -> vector<4xi1> + %ww = scf.if %valid_row_k -> (vector<4xf32>) { + %format44 = index.constant 44 : index + %packed = index.cmp eq, %weight_format, %format44 : index + %wh = scf.if %packed -> (vector<4xf32>) { + %v = func.call @ggml_q4k_native_row64_f32_vector4(%w_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %v = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %w_na, %row_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } + scf.yield %wh : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next0 = scf.if %valid0 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t0, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc0 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc0 : vector<4xf32> + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next1 = scf.if %valid1 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t1, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc1 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc1 : vector<4xf32> + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next2 = scf.if %valid2 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t2, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc2 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc2 : vector<4xf32> + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next3 = scf.if %valid3 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t3, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc3 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc3 : vector<4xf32> + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next4 = scf.if %valid4 -> (vector<4xf32>) { + %a = vector.load.mask %av[%t4, %block, %local_k], %input_mask, %zero : view<[%tokens]x[%blocks]x256xf32>, vector<4xi1>, vector<4xf32> + %r = vector.fmaf %ww, %a, %acc4 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc4 : vector<4xf32> + } + scf.yield %next0, %next1, %next2, %next3, %next4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %leader = index.cmp eq, %lane16, %c0 : index + %c5 = index.constant 5 : index + %stage_width = scf.select %partitioned, %c16, %c8 : index + %stage_elements = index.mul %stage_width, %c5 : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%stage_width]xf32> + %local_row0 = index.add %wave_add, %cohort : index + %partial_offset = index.mul %k_partition, %c4 : index + %local_row = index.add %local_row0, %partial_offset : index + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total0, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total1, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total2, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total3, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum8 = func.call @ggml_matmul_vector_reduce_f32(%total4, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %rows_per_group, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_group_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, 4), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %stage_width : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_origin = scf.if %partitioned -> (index) { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 64), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } else { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 36), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values0 = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %values = scf.if %partitioned -> (vector<4xf32>) { + %stage_origin1 = index.add %stage_origin, %c4 : index + %stage_origin2 = index.add %stage_origin1, %c4 : index + %stage_origin3 = index.add %stage_origin2, %c4 : index + %values1 = vector.load %flat_stage[%stage_origin1] : view<[%stage_elements]xf32> -> vector<4xf32> + %values2 = vector.load %flat_stage[%stage_origin2] : view<[%stage_elements]xf32> -> vector<4xf32> + %values3 = vector.load %flat_stage[%stage_origin3] : view<[%stage_elements]xf32> -> vector<4xf32> + %sum01 = vector.addf %values0, %values1 : vector<4xf32> + %sum012 = vector.addf %sum01, %values2 : vector<4xf32> + %sum = vector.addf %sum012, %values3 : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values0 : vector<4xf32> + } + %publish_format = config.get @ggml.matmul.vector.publish_format : index + template.apply<@ggml.matmul.vector_activation.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %publish_format, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %row_rounding = index.sub %rows_per_group, %c1 : index + %rounded_n = index.add %n, %row_rounding : index + %output_groups = index.div %rounded_n, %rows_per_group : index + template.apply<@ggml.matmul.vector_activation.finish_output_group>(%tokens, %n, %output_groups, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Copy the original Q8_1_x4 records for up to 2048 reduction elements. +// Call uniformly at a slab boundary; only the final slab may be partial. +func.def inline @ggml_matmul_vector_q8_1_x4_stage_slab(%input: buffer, %stage: buffer, %capacity: index, %slab_begin: index, %blocks: index, %global_row_bytes: offset) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c144 = index.constant 144 : index + %c576 = index.constant 576 : index + %c18 = index.constant 18 : index + %base = index.constant 0 : offset + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %workgroup_size = kernel.workgroup.size : index + %thread = kernel.workitem.id : index + %stage_packets = index.mul %capacity, %c144 : index + %stage_elements = index.mul %capacity, %c576 : index + %stage_view = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xi32> + %slab_limit = index.add %slab_begin, %c8 : index + %full_slab = index.cmp ule, %slab_limit, %blocks : index + %slab_end = scf.select %full_slab, %slab_limit, %blocks : index + %active_blocks0 = index.sub %slab_end, %slab_begin : index + %active_blocks = index.assume %active_blocks0 [range(%active_blocks0, 1, 8)] : index + %active_packets = index.mul %active_blocks, %c18 : index + %slab_byte_add = index.scale %slab_begin, %block_bytes : index, offset -> offset + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %copy_packet = [%thread to %stage_packets step %workgroup_size] { + %copy_token = index.div %copy_packet, %c144 : index + %copy_inner = index.rem %copy_packet, %c144 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + scf.if %copy_valid { + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %copy_bytes = index.scale %copy_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slab_byte_add : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + %copy_origin = index.mul %copy_packet, %c4 : index + vector.store %values, %stage_view[%copy_origin] : vector<4xi32>, view<[%stage_elements]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_matmul_vector_q8_publish(%tokens: index, %n: index, %rows_per_group: index, %row_base: index, %local_row: index, %block_lane: index, %channel_group: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %publish_format: index, %total_g0: f32, %total_g1: f32, %total_g2: f32, %total_g3: f32, %total_g4: f32, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %leader = index.cmp eq, %block_lane, %c0 : index + %width = scalar.constant 32 : i32 + %stage_rows = index.constant 5 : index + %stage_elements = index.mul %stage_rows, %rows_per_group : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%rows_per_group]xf32> + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum = scalar.addf %total_g0, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum = scalar.addf %total_g1, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum = scalar.addf %total_g2, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum = scalar.addf %total_g3, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum = scalar.addf %total_g4, %zero : f32 + %sum8 = func.call @ggml_matmul_vector_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%rows_per_group]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %rows_per_group, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %row_limit = index.sub %rows_per_group, %c4 : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, %row_limit), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %rows_per_group : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_limit = index.sub %stage_elements, %c4 : index + %stage_origin = index.assume %stage_origin0 [range(%stage_origin0, 0, %stage_limit), mul(%stage_origin0, 4)] : index + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %publish_format = config.get @ggml.matmul.vector.publish_format : index + template.apply<@ggml.matmul.vector_activation.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %publish_format, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %rounding = index.sub %rows_per_group, %c1 : index + %rounded_n = index.add %n, %rounding : index + %output_groups = index.div %rounded_n, %rows_per_group : index + template.apply<@ggml.matmul.vector_activation.finish_output_group>(%tokens, %n, %output_groups, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_q4_prefetch_native_words(%weight: buffer, %row_group_block: index, %block: index, %row_lane: index, %payload_field: index, %payload_word: index) -> (vector<4xi32>, vector<2xi32>) { + %zero = index.constant 0 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %group_block = index.add %row_group_block, %block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %header_view = buffer.view %weight[%wb] : buffer -> view<4xi32> + %code_view = buffer.view %weight[%wc] : buffer -> view<8x64x4xi32> + %headers = vector.load %header_view[%zero] : view<4xi32> -> vector<4xi32> + %raw = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + func.return %headers, %raw : vector<4xi32>, vector<2xi32> +} + +// Native Q4_K bytes and Q8_1_x4 activations, with one 16-lane cohort per row. +// Retain the wrapper publication hooks so bias, residual and RMSNorm fusions +// share this contraction without a separate dispatch or intermediate output. +func.def pure inline @ggml_matmul_vector_q4_q8_token_pairs(%capacity: index, %tokens: index, %k: index, %n: index) -> (i1) { + %c4 = index.constant 4 : index + %c4096 = index.constant 4096 : index + %capacity4 = index.cmp eq, %capacity, %c4 : index + %tokens4 = index.cmp eq, %tokens, %c4 : index + %four_tokens = scalar.andi %capacity4, %tokens4 : i1 + %long_input = index.cmp uge, %k, %c4096 : index + %wide_output = index.cmp uge, %n, %c4096 : index + %contracts = index.cmp ule, %n, %k : index + %wide = scalar.andi %long_input, %wide_output : i1 + %contracting = scalar.andi %wide, %contracts : i1 + %eligible = scalar.andi %four_tokens, %contracting : i1 + func.return %eligible : i1 +} + +func.def inline @ggml_matmul_vector_q4_q8_adjacent_scale_min(%scale0: i32, %scale1: i32, %scale2: i32, %pair0: index) -> (i32, i32, i32, i32) { + %pair = index.assume %pair0 [range(%pair0, 0, 3)] : index + %two = index.constant 2 : index + %sixteen = index.constant 16 : index + %low = index.cmp ult, %pair, %two : index + %pair_lane = index.rem %pair, %two : index + %shift_index = index.mul %pair_lane, %sixteen : index + %shift = index.cast %shift_index : index to i32 + %two_i32 = scalar.constant 2 : i32 + %four_i32 = scalar.constant 4 : i32 + %eight_i32 = scalar.constant 8 : i32 + %low_mask = scalar.constant 3855 : i32 + %high_mask = scalar.constant 12336 : i32 + %byte_mask = scalar.constant 255 : i32 + %high_shift = scalar.addi %shift, %two_i32 : i32 + %min_shift = scalar.addi %shift, %four_i32 : i32 + %s_source = scf.select %low, %scale0, %scale2 : i32 + %m_source = scf.select %low, %scale1, %scale2 : i32 + %s_high_shift = scf.select %low, %shift, %high_shift : i32 + %m_low_shift = scf.select %low, %shift, %min_shift : i32 + %s_lo_raw = scalar.shrui %s_source, %shift : i32 + %s_lo = scalar.andi %s_lo_raw, %low_mask : i32 + %s_hi_raw = scalar.shrui %scale0, %s_high_shift : i32 + %s_hi = scalar.andi %s_hi_raw, %high_mask : i32 + %s = scalar.ori %s_lo, %s_hi : i32 + %m_lo_raw = scalar.shrui %m_source, %m_low_shift : i32 + %m_lo = scalar.andi %m_lo_raw, %low_mask : i32 + %m_hi_raw = scalar.shrui %scale1, %s_high_shift : i32 + %m_hi = scalar.andi %m_hi_raw, %high_mask : i32 + %m = scalar.ori %m_lo, %m_hi : i32 + %s0 = scalar.andi %s, %byte_mask : i32 + %m0 = scalar.andi %m, %byte_mask : i32 + %s1 = scalar.shrui %s, %eight_i32 : i32 + %m1 = scalar.shrui %m, %eight_i32 : i32 + func.return %s0, %m0, %s1, %m1 : i32, i32, i32, i32 +} + +func.def inline @ggml_matmul_vector_q4_q8_dot8(%paired_tokens: i1, %weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) { + %value = scf.if %paired_tokens -> (f32) { + %zero = vector.constant 0 : vector<1xi32> + %wi0 = vector.extract %weight[0] : vector<2xi32> -> i32 + %ai0 = vector.extract %activation[0] : vector<2xi32> -> i32 + %ww0 = vector.splat %wi0 : vector<1xi32> + %aw0 = vector.splat %ai0 : vector<1xi32> + %w0 = vector.bitcast %ww0 : vector<1xi32> to vector<4xi8> + %a0 = vector.bitcast %aw0 : vector<1xi32> to vector<4xi8> + %d0 = vector.dot4i %w0, %a0, %zero : vector<4xi8>, vector<4xi8>, vector<1xi32> + %wi1 = vector.extract %weight[1] : vector<2xi32> -> i32 + %ai1 = vector.extract %activation[1] : vector<2xi32> -> i32 + %ww1 = vector.splat %wi1 : vector<1xi32> + %aw1 = vector.splat %ai1 : vector<1xi32> + %w1 = vector.bitcast %ww1 : vector<1xi32> to vector<4xi8> + %a1 = vector.bitcast %aw1 : vector<1xi32> to vector<4xi8> + %d1 = vector.dot4i %w1, %a1, %d0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %sum = vector.extract %d1[0] : vector<1xi32> -> i32 + %result = scalar.sitofp %sum : i32 to f32 + scf.yield %result : f32 + } else { + %original = func.call @ggml_dot_u8_s8_vector8_f32(%weight, %activation) : (vector<2xi32>, vector<2xi32>) -> (f32) + scf.yield %original : f32 + } + func.return %value : f32 +} + +func.def inline @ggml_matmul_vector_q4_q8_1_x4_activation_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_matmul_vector_q8_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %paired_tokens = func.call pure @ggml_matmul_vector_q4_q8_token_pairs(%capacity, %tokens, %k, %n) : (index, index, index, index) -> (i1) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %wg = index.add %channel_group, %c0 : index + %c64_rows = index.constant 64 : index + %direct_rows = index.div %c8, %k_partitions : index + %rows_per_group = scf.select %staged, %c64_rows, %direct_rows : index + %row_base = index.mul %wg, %rows_per_group : index + %wave_rows = index.div %c4, %k_partitions : index + %wave_row = index.mul %wave, %wave_rows : index + %cohort_width = index.mul %c16, %k_partitions : index + %row_half = index.div %lane, %cohort_width : index + %cohort_lane = index.rem %lane, %cohort_width : index + %local_row = index.add %wave_row, %row_half : index + %row = index.add %row_base, %local_row : index + %valid_row = index.cmp ult, %row, %n : index + %cohort = index.div %lane, %c16 : index + %block_lane = index.rem %lane, %c16 : index + %pair = index.div %block_lane, %c4 : index + %packet = index.rem %block_lane, %c4 : index + %group0 = index.mul %pair, %c2 : index + %group1 = index.add %group0, %c1 : index + %word_page = index.mul %pair, %c8 : index + %word_add = index.mul %packet, %c2 : index + %word_index = index.add %word_page, %word_add : index + %meta_lane = index.cmp eq, %packet, %c0 : index + %blocks = index.div %k, %c256 : index + %partition_blocks = index.div %blocks, %k_partitions : index + %k_team = index.div %cohort_lane, %c16 : index + %k_begin = index.mul %k_team, %partition_blocks : index + %weight_row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %weight_row_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_groups = index.div %k, %c128 : index + %global_q8_row_bytes = index.scale %q8_groups, %block_bytes : index, offset -> offset + %slab_row_bytes = index.constant 2304 : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %q8_half = index.div %pair, %c2 : index + %q8_pair = index.rem %pair, %c2 : index + %q8_group0 = index.mul %q8_pair, %c2 : index + %q8_group1 = index.add %q8_group0, %c1 : index + %q8_words_base = index.mul %q8_group0, %c8 : index + %q8_word0 = index.add %q8_words_base, %word_add : index + %q8_word1 = index.add %q8_word0, %c8 : index + %q8_meta0 = index.mul %q8_group0, %c2 : index + %q8_meta1 = index.mul %q8_group1, %c2 : index + %q8_sum0 = index.add %q8_meta0, %c1 : index + %q8_sum1 = index.add %q8_meta1, %c1 : index + %c64 = index.constant 64 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %channel_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %payload_field = index.div %word_index, %c4 : index + %payload_word = index.rem %word_index, %c4 : index + %input_na, %gate_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = index.constant 8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_header = vector.constant 0 : vector<4xi32> + %zero_raw = vector.constant 0 : vector<2xi32> + %prime_header, %prime_raw = scf.if %staged -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %c0, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %zero_header, %zero_raw : vector<4xi32>, vector<2xi32> + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_header, %last_raw = scf.for %block_base = [%c0 to %partition_blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_header = %prime_header : vector<4xi32>, %carry_raw = %prime_raw : vector<2xi32>) -> (f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32>) { + %block_in_slab = index.rem %block_base, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %new_slab = scalar.andi %staged, %slab_boundary : i1 + scf.if %new_slab { + func.call @ggml_matmul_vector_q8_1_x4_stage_slab(%input_na, %q8_stage, %capacity, %block_base, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %block = index.add %block_base, %k_begin : index + %valid_block = index.cmp ult, %block, %blocks : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %group_block = index.add %row_group_block, %safe_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %g_header_view = buffer.view %gate_na[%wb] : buffer -> view<4xi32> + %g_code_view = buffer.view %gate_na[%wc] : buffer -> view<8x64x4xi32> + %g_headers = scf.if %staged -> (vector<4xi32>) { + scf.yield %carry_header : vector<4xi32> + } else { + %header = vector.load %g_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %header : vector<4xi32> + } + %g_dm_i32 = vector.extract %g_headers[0] : vector<4xi32> -> i32 + %g_dm_word = vector.splat %g_dm_i32 : vector<1xi32> + %g_dm = vector.bitcast %g_dm_word : vector<1xi32> to vector<2xf16> + %g_d0h = vector.extract %g_dm[0] : vector<2xf16> -> f16 + %g_d1h = vector.extract %g_dm[1] : vector<2xf16> -> f16 + %g_d0 = scalar.extf %g_d0h : f16 to f32 + %g_d1 = scalar.extf %g_d1h : f16 to f32 + %g_s0 = vector.extract %g_headers[1] : vector<4xi32> -> i32 + %g_s1 = vector.extract %g_headers[2] : vector<4xi32> -> i32 + %g_s2 = vector.extract %g_headers[3] : vector<4xi32> -> i32 + %pair_scale0, %pair_min0, %pair_scale1, %pair_min1 = scf.if %paired_tokens -> (i32, i32, i32, i32) { + %s0, %m0, %s1, %m1 = func.call @ggml_matmul_vector_q4_q8_adjacent_scale_min(%g_s0, %g_s1, %g_s2, %pair) : (i32, i32, i32, index) -> (i32, i32, i32, i32) + scf.yield %s0, %m0, %s1, %m1 : i32, i32, i32, i32 + } else { + %zero_i32 = scalar.constant 0 : i32 + scf.yield %zero_i32, %zero_i32, %zero_i32, %zero_i32 : i32, i32, i32, i32 + } + %g_scale0i, %g_min0i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale0, %pair_min0 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale0f = scalar.uitofp %g_scale0i : i32 to f32 + %g_min0f = scalar.uitofp %g_min0i : i32 to f32 + %g_ws0 = scalar.mulf %g_d0, %g_scale0f : f32 + %g_wm0 = scalar.mulf %g_d1, %g_min0f : f32 + %g_scale1i, %g_min1i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale1, %pair_min1 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale1f = scalar.uitofp %g_scale1i : i32 to f32 + %g_min1f = scalar.uitofp %g_min1i : i32 to f32 + %g_ws1 = scalar.mulf %g_d0, %g_scale1f : f32 + %g_wm1 = scalar.mulf %g_d1, %g_min1f : f32 + %g_raw = scf.if %staged -> (vector<2xi32>) { + scf.yield %carry_raw : vector<2xi32> + } else { + %raw = vector.load %g_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %raw : vector<2xi32> + } + %g_mask = vector.constant 252645135 : vector<2xi32> + %g_shift = vector.constant 4 : vector<2xi32> + %g_lo = vector.andi %g_raw, %g_mask : vector<2xi32> + %g_high = vector.shrui %g_raw, %g_shift : vector<2xi32> + %g_hi = vector.andi %g_high, %g_mask : vector<2xi32> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_header, %next_raw = scf.if %has_next -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %next_block, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %carry_header, %carry_raw : vector<4xi32>, vector<2xi32> + } + %input_block = scf.select %staged, %block_in_slab, %safe_block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %q8_half : index + %q8_block_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next_g0 = scf.if %valid0 -> (f32) { + %tb = index.scale %t0, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ad0h = view.load %ds_view[%q8_meta0] : view<8xf16> -> f16 + %ad1h = view.load %ds_view[%q8_meta1] : view<8xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = view.load %ds_view[%q8_sum0] : view<8xf16> -> f16 + %as1h = view.load %ds_view[%q8_sum1] : view<8xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g0, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next_g1 = scf.if %valid1 -> (f32) { + %tb = index.scale %t1, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g1, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + scf.if %paired_tokens { + scf.schedule.fence + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next_g2 = scf.if %valid2 -> (f32) { + %tb = index.scale %t2, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g2, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next_g3 = scf.if %valid3 -> (f32) { + %tb = index.scale %t3, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g3, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next_g4 = scf.if %valid4 -> (f32) { + %tb = index.scale %t4, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_matmul_vector_q4_q8_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g4, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_header, %next_raw : f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32> + } + %combined0 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g0, %k_partitions) : (f32, index) -> (f32) + %combined1 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g1, %k_partitions) : (f32, index) -> (f32) + %combined2 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g2, %k_partitions) : (f32, index) -> (f32) + %combined3 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g3, %k_partitions) : (f32, index) -> (f32) + %combined4 = func.call @ggml_matmul_vector_combine_k_partitions(%total_g4, %k_partitions) : (f32, index) -> (f32) + %publish_format = config.get @ggml.matmul.vector.publish_format : index + func.call @ggml_matmul_vector_q8_publish(%tokens, %n, %rows_per_group, %row_base, %local_row, %cohort_lane, %channel_group, %output_accumulation, %output_unary_op, %epsilon, %publish_format, %combined0, %combined1, %combined2, %combined3, %combined4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, index, index, f32, index, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Native Q6_K carrier with signed I8 dot products and F32 scaling. +func.def inline @ggml_matmul_vector_q6_prefetch_native_words(%weight_na: buffer, %row_group_block: index, %block: index, %row_lane: index, %half16: index, %group: index, %half128: index, %ql_side: index, %packet: index) -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %c2 = index.constant 2 : index + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %group_block = index.add %row_group_block, %block : index + %block_base = index.scale %group_block, %record_bytes : index, offset -> offset + %scale_base = index.add %block_base, %scale_add : offset + %raw_base = index.add %block_base, %raw_add : offset + %d_view = buffer.view %weight_na[%block_base] : buffer -> view<64xf16> + %s_view = buffer.view %weight_na[%scale_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_na[%raw_base] : buffer -> view<2x3x64x8xi32> + %dh = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %si = view.load %s_view[%row_lane, %half16, %group] : view<64x2x8xi8> -> i8 + %ql = vector.load %raw_view[%half128, %ql_side, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + %qh = vector.load %raw_view[%half128, %c2, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + func.return %dh, %si, %ql, %qh : f16, i8, vector<4xi32>, vector<4xi32> +} + +func.def inline @ggml_matmul_vector_q6_q8_activation_values(%pipelined: i1, %token: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index) -> (vector<16xi8>, f32) { + %code_offset = index.constant 16 : offset + %a, %ad = scf.if %pipelined -> (vector<16xi8>, f32) { + %base = index.constant 0 : offset + %c4 = index.constant 4 : index + %q8_block = index.assume %q8_block0 [range(%q8_block0, 0, 7)] : index + %word = index.add %q8_word0, %c4 : index + %q_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x36xi32> + %ds_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x72xf16> + %words = vector.load %q_view[%current_slot, %token, %q8_block, %word] : view<2x[%capacity]x8x36xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%current_slot, %token, %q8_block, %q8_meta] : view<2x[%capacity]x8x72xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } else { + %tb = index.scale %token, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %words = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%q8_meta] : view<8xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } + func.return %a, %ad : vector<16xi8>, f32 +} + +func.def inline @ggml_matmul_vector_q6_q8_dot16(%pipelined: i1, %w: vector<16xi8>, %a: vector<16xi8>) -> (i32) { + %sum = scf.if %pipelined -> (i32) { + %di0 = vector.constant 0 : vector<1xi32> + %dw0 = vector.slice %w[0] : vector<16xi8> -> vector<4xi8> + %da0 = vector.slice %a[0] : vector<16xi8> -> vector<4xi8> + %di1 = vector.dot4i %dw0, %da0, %di0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw1 = vector.slice %w[4] : vector<16xi8> -> vector<4xi8> + %da1 = vector.slice %a[4] : vector<16xi8> -> vector<4xi8> + %di2 = vector.dot4i %dw1, %da1, %di1 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw2 = vector.slice %w[8] : vector<16xi8> -> vector<4xi8> + %da2 = vector.slice %a[8] : vector<16xi8> -> vector<4xi8> + %di3 = vector.dot4i %dw2, %da2, %di2 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw3 = vector.slice %w[12] : vector<16xi8> -> vector<4xi8> + %da3 = vector.slice %a[12] : vector<16xi8> -> vector<4xi8> + %di4 = vector.dot4i %dw3, %da3, %di3 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %chain_sum = vector.extract %di4[0] : vector<1xi32> -> i32 + scf.yield %chain_sum : i32 + } else { + %init = vector.constant 0 : vector<4xi32> + %dot = vector.dot4i %w, %a, %init : vector<16xi8>, vector<16xi8>, vector<4xi32> + %seed = scalar.constant 0 : i32 + %direct_sum = vector.reduce %dot, %seed : vector<4xi32>, i32 + scf.yield %direct_sum : i32 + } + func.return %sum : i32 +} + +func.def inline @ggml_matmul_vector_q8_1_x4_prefetch_half_slab(%input: buffer, %capacity0: index, %slot_begin: index, %blocks: index, %global_row_bytes: offset) -> (vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c18 = index.constant 18 : index + %c72 = index.constant 72 : index + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %copy_packet = index.rem %thread, %total_packets : index + %copy_token = index.div %copy_packet, %c72 : index + %copy_inner = index.rem %copy_packet, %c72 : index + %remaining = index.sub %blocks, %slot_begin : index + %full = index.cmp uge, %remaining, %c4 : index + %active0 = scf.select %full, %c4, %remaining : index + %active = index.assume %active0 [range(%active0, 1, 4)] : index + %active_packets = index.mul %active, %c18 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + %safe_inner = scf.select %copy_valid, %copy_inner, %c0 : index + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %slot_bytes = index.scale %slot_begin, %block_bytes : index, offset -> offset + %copy_bytes = index.scale %safe_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slot_bytes : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + func.return %values : vector<4xi32> +} +func.def inline @ggml_matmul_vector_q8_1_x4_publish_half_slab(%stage: buffer, %capacity0: index, %slot0: index, %values: vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %slot = index.assume %slot0 [range(%slot0, 0, 1)] : index + %c4 = index.constant 4 : index + %c72 = index.constant 72 : index + %base = index.constant 0 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %thread_valid = index.cmp ult, %thread, %total_packets : index + scf.if %thread_valid { + %copy_token0 = index.div %thread, %c72 : index + %copy_token = index.assume %copy_token0 [range(%copy_token0, 0, 4), lt(%copy_token0, %capacity)] : index + %copy_inner = index.rem %thread, %c72 : index + %origin = index.mul %copy_inner, %c4 : index + %destination = buffer.view %stage[%base] : buffer -> view<2x[%capacity]x288xi32> + vector.store %values, %destination[%slot, %copy_token, %origin] : vector<4xi32>, view<2x[%capacity]x288xi32> + } + func.return +} + + +func.def inline @ggml_matmul_vector_q6_q8_accumulate(%pipelined: i1, %tokens: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index, %w: vector<16xi8>, %ws: f32, %acc_g0: f32, %acc_g1: f32, %acc_g2: f32, %acc_g3: f32, %acc_g4: f32) -> (f32, f32, f32, f32, f32) { + %code_offset = index.constant 16 : offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %computed_g0 = scf.if %valid0 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t0, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g0 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %computed_g1 = scf.if %valid1 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t1, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g1 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %computed_g2 = scf.if %valid2 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t2, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g2 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %computed_g3 = scf.if %valid3 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t3, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g3 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %computed_g4 = scf.if %valid4 -> (f32) { + %a, %ad = func.call @ggml_matmul_vector_q6_q8_activation_values(%pipelined, %t4, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_matmul_vector_q6_q8_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g4 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + func.return %computed_g0, %computed_g1, %computed_g2, %computed_g3, %computed_g4 : f32, f32, f32, f32, f32 +} + +func.def inline @ggml_matmul_vector_q6_q8_1_x4_activation_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_matmul_vector_q8_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %rows_per_group = scf.select %staged, %c32, %c8 : index + %row_base = index.mul %channel_group, %rows_per_group : index + %wave_row = index.mul %wave, %c4 : index + %cohort = index.div %lane, %c16 : index + %local_row = index.add %wave_row, %cohort : index + %row = index.add %row_base, %local_row : index + %block_lane = index.rem %lane, %c16 : index + %group = index.div %block_lane, %c2 : index + %half16 = index.rem %block_lane, %c2 : index + %half128 = index.div %group, %c4 : index + %group_in_half = index.rem %group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift0 = index.mul %nibble, %c4 : index + %nibble_shift1 = index.cast %nibble_shift0 : index to i32 + %nibble_shift = vector.splat %nibble_shift1 : vector<4xi32> + %high_shift0 = index.mul %group_in_half, %c2 : index + %high_shift1 = index.cast %high_shift0 : index to i32 + %high_shift = vector.splat %high_shift1 : vector<4xi32> + %packet = index.mul %half16, %c4 : index + %q8_group_words = index.mul %group_in_half, %c8 : index + %q8_word0 = index.add %q8_group_words, %packet : index + %q8_meta = index.mul %group_in_half, %c2 : index + %blocks = index.div %k, %c256 : index + %q8_groups = index.div %k, %c128 : index + %q8_block_bytes = index.constant 144 : offset + %global_q8_row_bytes = index.scale %q8_groups, %q8_block_bytes : index, offset -> offset + %native_block_bytes = index.constant 210 : index + %cache_bytes = index.constant 33554432 : index + %weight_blocks = index.mul %blocks, %n : index + %weight_bytes = index.mul %weight_blocks, %native_block_bytes : index + %exceeds_cache = index.cmp uge, %weight_bytes, %cache_bytes : index + %pipelined = scalar.andi %staged, %exceeds_cache : i1 + %full_slab_row_bytes = index.constant 2304 : offset + %half_slab_row_bytes = index.constant 1152 : offset + %slab_row_bytes = scf.select %pipelined, %half_slab_row_bytes, %full_slab_row_bytes : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %code_offset = index.constant 16 : offset + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %staged_row_group = index.div %channel_group, %c2 : index + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %staged_row_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %input_na, %weight_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = scf.select %pipelined, %c4, %c8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_half = scalar.constant 0.0 : f16 + %zero_code = scalar.constant 0 : i8 + %zero_words = vector.constant 0 : vector<4xi32> + %prime_d, %prime_s, %prime_ql, %prime_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %pd, %ps, %pql, %pqh = func.call @ggml_matmul_vector_q6_prefetch_native_words(%weight_na, %row_group_block, %c0, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %pd, %ps, %pql, %pqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %zero_half, %zero_code, %zero_words, %zero_words : f16, i8, vector<4xi32>, vector<4xi32> + } + scf.if %pipelined { + %prime_q8 = func.call @ggml_matmul_vector_q8_1_x4_prefetch_half_slab(%input_na, %capacity, %c0, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + func.call @ggml_matmul_vector_q8_1_x4_publish_half_slab(%q8_stage, %capacity, %c0, %prime_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_d, %last_s, %last_ql, %last_qh = scf.for %block = [%c0 to %blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_d = %prime_d : f16, %carry_s = %prime_s : i8, %carry_ql = %prime_ql : vector<4xi32>, %carry_qh = %prime_qh : vector<4xi32>) -> (f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32>) { + %block_in_slab = index.rem %block, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %direct_stage = scalar.andi %staged, %slab_boundary : i1 + %true = scalar.constant true : i1 + %not_pipelined = scalar.xori %pipelined, %true : i1 + %stage_current = scalar.andi %direct_stage, %not_pipelined : i1 + scf.if %stage_current { + func.call @ggml_matmul_vector_q8_1_x4_stage_slab(%input_na, %q8_stage, %capacity, %block, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %slab_number = index.div %block, %slab_blocks : index + %current_slot = index.rem %slab_number, %c2 : index + %following_slab = index.add %slab_number, %c1 : index + %following_begin = index.mul %following_slab, %slab_blocks : index + %following_exists = index.cmp ult, %following_begin, %blocks : index + %next_slot = index.rem %following_slab, %c2 : index + %staged_future = scalar.andi %pipelined, %following_exists : i1 + %current_d, %current_s, %current_ql, %current_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + %wd, %ws8, %wql, %wqh = func.call @ggml_matmul_vector_q6_prefetch_native_words(%weight_na, %row_group_block, %block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } + %d = scalar.extf %current_d : f16 to f32 + %sf = scalar.sitofp %current_s : i8 to f32 + %ws = scalar.mulf %d, %sf : f32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %four = vector.constant 4 : vector<4xi32> + %ql_shifted = vector.shrui %current_ql, %nibble_shift : vector<4xi32> + %low = vector.andi %ql_shifted, %nibble_mask : vector<4xi32> + %qh_shifted = vector.shrui %current_qh, %high_shift : vector<4xi32> + %high0 = vector.andi %qh_shifted, %high_mask : vector<4xi32> + %high = vector.shli %high0, %four : vector<4xi32> + %code = vector.ori %low, %high : vector<4xi32> + // Guard bits prevent subtraction from borrowing across packed bytes. + %guard_mask = vector.constant -2139062144 : vector<4xi32> + %center = vector.constant 538976288 : vector<4xi32> + %guarded = vector.ori %code, %guard_mask : vector<4xi32> + %centered = vector.subi %guarded, %center : vector<4xi32> + %signed_code = vector.xori %centered, %guard_mask : vector<4xi32> + %w = vector.bitcast %signed_code : vector<4xi32> to vector<16xi8> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_d, %next_s, %next_ql, %next_qh = scf.if %has_next -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %wd, %ws8, %wql, %wqh = func.call @ggml_matmul_vector_q6_prefetch_native_words(%weight_na, %row_group_block, %next_block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } + %input_block = scf.select %staged, %block_in_slab, %block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %half128 : index + %q8_local_block_add = index.scale %q8_block, %q8_block_bytes : index, offset -> offset + %slot_row = index.mul %current_slot, %capacity : index + %slot_byte_add0 = index.scale %slot_row, %slab_row_bytes : index, offset -> offset + %slot_byte_add = scf.select %pipelined, %slot_byte_add0, %base : offset + %q8_block_add = index.add %slot_byte_add, %q8_local_block_add : offset + %slot_complete = index.cmp eq, %block_in_slab, %c3 : index + %publish_next = scalar.andi %staged_future, %slot_complete : i1 + %next_g0, %next_g1, %next_g2, %next_g3, %next_g4 = scf.if %publish_next -> (f32, f32, f32, f32, f32) { + %next_q8 = func.call @ggml_matmul_vector_q8_1_x4_prefetch_half_slab(%input_na, %capacity, %following_begin, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_matmul_vector_q6_q8_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + func.call @ggml_matmul_vector_q8_1_x4_publish_half_slab(%q8_stage, %capacity, %next_slot, %next_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } else { + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_matmul_vector_q6_q8_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_d, %next_s, %next_ql, %next_qh : f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32> + } + %publish_format = config.get @ggml.matmul.vector.publish_format : index + func.call @ggml_matmul_vector_q8_publish(%tokens, %n, %rows_per_group, %row_base, %local_row, %block_lane, %channel_group, %output_accumulation, %output_unary_op, %epsilon, %publish_format, %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, index, index, f32, index, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + + +template.def<@ggml.matmul.vector_activation.core> device @ggml_matmul_vector_activation_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %activation_format = config.get @ggml.matmul.vector.activation_format : index + %q8_format = index.constant 9 : index + %uses_q8 = index.cmp eq, %activation_format, %q8_format : index + scf.if %uses_q8 { + %q6_format = index.constant 46 : index + %uses_q6 = index.cmp eq, %weight_format, %q6_format : index + scf.if %uses_q6 { + func.call @ggml_matmul_vector_q6_q8_1_x4_activation_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_matmul_vector_q4_q8_1_x4_activation_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } else { + %token_capacity = config.get @ggml.workload.token_capacity : index + %publish_format = config.get @ggml.matmul.vector.publish_format : index + %pair_decode = func.call pure @ggml_matmul_vector_uses_f32_decode_pair(%token_capacity, %input_size0, %publish_format) : (index, index, index) -> (i1) + scf.if %pair_decode { + func.call @ggml_matmul_vector_f32_decode_pair(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %publish_format, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_matmul_vector_f32_activation_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_group, %token_group, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_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/matmul/vector_decode_pair.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_decode_pair.loom new file mode 100644 index 000000000000..897e660bfc97 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_decode_pair.loom @@ -0,0 +1,225 @@ +func.decl @ggml_q1_0_load_decode_state(%weight: buffer, %row_byte_base: offset, %k: index) -> (vector<1xf16>) +func.decl @ggml_q1_0_preloaded_f32_vector4(%weight: buffer, %row_byte_base: offset, %k: index, %d_vector: vector<1xf16>) -> (vector<4xf32>) +template.decl @ggml.matmul.vector_activation.publish_pair2(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: i1, %arg6: f32, %arg7: f32, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +func.def inline @ggml_matmul_vector_decode_pair_load_f32_block(%token_count: index, %input_size: index, %token: index, %block: index, %lane: index, %input: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %padded_input_size = index.add %input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block0 = scf.select %valid_block, %block, %c0 : index + %safe_block, %launch_block_count = index.assume %safe_block0, %block_count [lt(%safe_block0, %block_count)] : index, index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %vector_index = index.rem %itid, %c8 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%launch_block_count]x256xf32> + %vector_half_base = index.mul %vector_half, %c128 : index + %vector_offset = index.mul %vector_index, %c4 : index + %input_index0 = index.add %vector_half_base, %vector_offset : index + %input_index1 = index.add %input_index0, %c32 : index + %input_index2 = index.add %input_index0, %c64 : index + %input_index3 = index.add %input_index0, %c96 : index + %block_k_base = index.mul %safe_block, %c256 : index + %k0 = index.add %block_k_base, %input_index0 : index + %k1 = index.add %block_k_base, %input_index1 : index + %k2 = index.add %block_k_base, %input_index2 : index + %k3 = index.add %block_k_base, %input_index3 : index + %mask0 = vector.mask.range [%k0 to %input_size step %c1] : index -> vector<4xi1> + %mask1 = vector.mask.range [%k1 to %input_size step %c1] : index -> vector<4xi1> + %mask2 = vector.mask.range [%k2 to %input_size step %c1] : index -> vector<4xi1> + %mask3 = vector.mask.range [%k3 to %input_size step %c1] : index -> vector<4xi1> + %input0 = vector.load.mask %input_view[%token, %safe_block, %input_index0], %mask0, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input1 = vector.load.mask %input_view[%token, %safe_block, %input_index1], %mask1, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input2 = vector.load.mask %input_view[%token, %safe_block, %input_index2], %mask2, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input3 = vector.load.mask %input_view[%token, %safe_block, %input_index3], %mask3, %zero : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> +} + +func.def inline @ggml_matmul_vector_decode_pair_block_row(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %input_size: index, %row: index, %block: index, %lane: index, %weight: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero = vector.constant 0.0 : vector<4xf32> + %zero_f32 = scalar.constant 0.0 : f32 + %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index + %padded_input_size = index.add %input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %valid_block = index.cmp ult, %block, %block_count : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %itid = index.rem %bounded_lane, %c16 : index + %vector_half = index.div %itid, %c8 : index + %packet = index.rem %itid, %c8 : index + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %group_base = index.mul %vector_half, %c4 : index + %group0 = index.add %group_base, %c0 : index + %group1 = index.add %group_base, %c1 : index + %group2 = index.add %group_base, %c2 : index + %group3 = index.add %group_base, %c3 : index + %block_k_base = index.mul %safe_block, %c256 : index + %half_base = index.mul %vector_half, %c128 : index + %packet_k = index.mul %packet, %c4 : index + %k_half = index.add %block_k_base, %half_base : index + %k0 = index.add %k_half, %packet_k : index + %k1 = index.add %k0, %c32 : index + %k2 = index.add %k0, %c64 : index + %k3 = index.add %k0, %c96 : index + %valid_k0 = index.cmp ult, %k0, %input_size : index + %valid_k1 = index.cmp ult, %k1, %input_size : index + %valid_k2 = index.cmp ult, %k2, %input_size : index + %valid_k3 = index.cmp ult, %k3, %input_size : index + %metadata_format = index.constant 10 : index + %metadata_enabled = index.cmp eq, %weight_format, %metadata_format : index + %metadata_block_safe = scalar.andi %metadata_enabled, %valid_block : i1 + %metadata_safe = scalar.andi %metadata_block_safe, %valid_k0 : i1 + %cached0 = scf.if %metadata_safe -> (vector<1xf16>) { + %cached0 = func.call @ggml_q1_0_load_decode_state(%weight, %row_base, %k0) : (buffer, offset, index) -> (vector<1xf16>) + scf.yield %cached0 : vector<1xf16> + } else { + %cached0 = vector.constant 0.0 : vector<1xf16> + scf.yield %cached0 : vector<1xf16> + } + %weight0 = scf.if %valid_k0 -> (vector<4xf32>) { + %decoded = scf.if %metadata_enabled -> (vector<4xf32>) { + %v = func.call @ggml_q1_0_preloaded_f32_vector4(%weight, %row_base, %k0, %cached0) : (buffer, offset, index, vector<1xf16>) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %fallback = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group0, %packet, %k0) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %fallback : vector<4xf32> + } + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight1 = scf.if %valid_k1 -> (vector<4xf32>) { + %decoded = scf.if %metadata_enabled -> (vector<4xf32>) { + %v = func.call @ggml_q1_0_preloaded_f32_vector4(%weight, %row_base, %k1, %cached0) : (buffer, offset, index, vector<1xf16>) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %fallback = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group1, %packet, %k1) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %fallback : vector<4xf32> + } + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight2 = scf.if %valid_k2 -> (vector<4xf32>) { + %decoded = scf.if %metadata_enabled -> (vector<4xf32>) { + %v = func.call @ggml_q1_0_preloaded_f32_vector4(%weight, %row_base, %k2, %cached0) : (buffer, offset, index, vector<1xf16>) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %fallback = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group2, %packet, %k2) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %fallback : vector<4xf32> + } + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %weight3 = scf.if %valid_k3 -> (vector<4xf32>) { + %decoded = scf.if %metadata_enabled -> (vector<4xf32>) { + %v = func.call @ggml_q1_0_preloaded_f32_vector4(%weight, %row_base, %k3, %cached0) : (buffer, offset, index, vector<1xf16>) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %fallback = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_base, %input_size, %safe_block, %group3, %packet, %k3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %fallback : vector<4xf32> + } + scf.yield %decoded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %mul0 = vector.mulf %input0, %weight0 : vector<4xf32> + %mul1 = vector.mulf %input1, %weight1 : vector<4xf32> + %mul2 = vector.mulf %input2, %weight2 : vector<4xf32> + %mul3 = vector.mulf %input3, %weight3 : vector<4xf32> + %sum01 = vector.addf %mul0, %mul1 : vector<4xf32> + %sum23 = vector.addf %mul2, %mul3 : vector<4xf32> + %sum = vector.addf %sum01, %sum23 : vector<4xf32> + %block_contribution = vector.reduce %sum, %zero_f32 : vector<4xf32>, f32 + %contribution = scf.select %valid_block, %block_contribution, %zero_f32 : f32 + func.return %contribution : f32 +} + +func.def inline @ggml_matmul_vector_f32_decode_pair(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %tokens = index.assume %token_count [range(%token_count, 1, 1)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %lane0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero = scalar.constant 0.0 : f32 + %token_valid = index.cmp ult, %token_group, %tokens : index + %safe_token0 = scf.select %token_valid, %token_group, %c0 : index + %token, %launch_tokens = index.assume %safe_token0, %tokens [lt(%safe_token0, %tokens)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.mul %channel_group, %c2 : index + %row_valid = index.cmp ult, %row00, %n : index + %safe_row0 = scf.select %row_valid, %row00, %c0 : index + %row0, %launch_n = index.assume %safe_row0, %n [lt(%safe_row0, %n)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %launch_n : index + %publish_output = scalar.andi %token_valid, %row_valid : i1 + %padded_k = index.add %k, %c255 : index + %block_count = index.div %padded_k, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %zero : f32, %row_acc1 = %zero : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + %input0, %input1, %input2, %input3 = func.call @ggml_matmul_vector_decode_pair_load_f32_block(%launch_tokens, %k, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_matmul_vector_decode_pair_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %k, %row0, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + %row1_contribution = func.call @ggml_matmul_vector_decode_pair_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %k, %row1, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %zero : f32 + } + %next0 = scalar.addf %row_acc0, %contribution0 : f32 + %next1 = scalar.addf %row_acc1, %contribution1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %sum0 = kernel.workgroup.reduce %acc0 : f32 + %sum1 = kernel.workgroup.reduce %acc1 : f32 + %is_lane_zero = index.cmp eq, %lane, %c0 : index + %publish_pair = scalar.andi %publish_output, %is_lane_zero : i1 + template.apply<@ggml.matmul.vector_activation.publish_pair2>(%publish_pair, %token, %row0, %launch_tokens, %launch_n, %row1_valid, %sum0, %sum1, %output_accumulation, %output_unary_op, %publish_format, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters, %next_output) : (i1, index, index, index, index, i1, f32, f32, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + func.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom new file mode 100644 index 000000000000..5b782b08a182 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_postops.loom @@ -0,0 +1,46 @@ +template.decl @ggml.matmul.vector_activation.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.vector_activation.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.vector_activation.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul.vector_activation.add_bias_vector4> device @ggml_matmul_vector_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.vector_activation.add_residual_vector4> device @ggml_matmul_vector_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul.vector_activation.store_vector4> device @ggml_matmul_vector_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom new file mode 100644 index 000000000000..6008d253a25d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish.loom @@ -0,0 +1,21 @@ +// Shared publication helpers for vector-shaped matmul outputs. +template.decl @ggml.matmul.vector_publish.next_vector4(%publish_word: i1, %publish_format: index, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %next_output: buffer) + +template.def<@ggml.matmul.vector_publish.next_vector4> device @ggml_matmul_vector_publish_next_vector4(%publish_word: i1, %publish_format: index, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %f16_format = index.constant 1 : index + %base = index.constant 0 : offset + %publishes_f16 = index.cmp eq, %publish_format, %f16_format : index + %publishes_next = scalar.andi %publish_word, %publishes_f16 : i1 + scf.if %publishes_next { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %next_output_noalias = buffer.assume.noalias %next_output : buffer + %next_output_view = buffer.view %next_output_noalias[%base] : buffer -> view<[%token_count]x[%output_size]xf16> + %f16_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store.mask %f16_values, %next_output_view[%token, %channel], %mask : vector<4xf16>, view<[%token_count]x[%output_size]xf16>, vector<4xi1> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom new file mode 100644 index 000000000000..64738a8ff0a8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_publish_f32.loom @@ -0,0 +1,22 @@ +// F32 publication stages shared across vector implementations. +template.decl @ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32(%selected_count: index, %output_size: index, %selected_outputs: buffer, %values: buffer, %output: buffer) + +template.def<@ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32> device @ggml_matmul_vector_publish_f32_scatter_selected_outputs_wave32(%selected_count: index, %output_size: index, %selected_outputs: buffer, %values: buffer, %output: buffer) { + %base = index.constant 0 : offset + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index + %publishes = index.cmp ult, %tid, %selected_count : index + %selected_noalias, %values_noalias, %output_noalias = buffer.assume.noalias %selected_outputs, %values, %output : buffer, buffer, buffer + %selected_view = buffer.view %selected_noalias[%base] : buffer -> view<64xi32> + %values_view = buffer.view %values_noalias[%base] : buffer -> view<64xf32> + scf.if %publishes { + %output_i32 = view.load %selected_view[%tid] : view<64xi32> -> i32 + %output0 = index.cast %output_i32 : i32 to index + %output1 = index.assume %output0 [range(%output0, 0, 262143)] : index + %selected_output, %bounded_output_size = index.assume %output1, %output_size [lt(%output1, %output_size)] : index, index + %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%bounded_output_size]xf32> + %value = view.load %values_view[%tid] : view<64xf32> -> f32 + view.store %value, %output_view[%selected_output] : f32, view<[%bounded_output_size]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom new file mode 100644 index 000000000000..93d341238477 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/vector_q6_f32_body.loom @@ -0,0 +1,748 @@ +// Shared Q6-weight vector matmul reductions. +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector.weight_offset : %value: index where [range(%value, 0, 2147483647), mul(%value, 256)] + +template.decl @ggml.matmul.vector_q6_f32.store_fragment_outputs_wave32(%result: vector<16xf16>, %token_offset: index, %token_start: index, %bounded_token_count: index, %bounded_output_size: index, %channel_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) + +template.decl @ggml.matmul.vector_q6_f32.store_dense_outputs_wave32(%result: vector<8xf32>, %token_offset: index, %token_start: index, %token_count: index, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer, %output_accumulation: index) + +template.decl @ggml.matmul.vector_q6_f32.store_single_row_overwrite_wave32(%result: vector<8xf32>, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_packed_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_i8_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_selected_outputs(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %selected_count0: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %output: buffer) + +template.def<@ggml.matmul.vector_q6_f32.store_fragment_outputs_wave32> device @ggml_matmul_vector_q6_f32_store_fragment_outputs_wave32(%result: vector<16xf16>, %token_offset: index, %token_start: index, %bounded_token_count: index, %bounded_output_size: index, %channel_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) { + %lane = kernel.subgroup.lane.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c0_offset = index.constant 0 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %publish_token0 = index.div %lane, %c2 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c2 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index + %publish_channel_add = index.mul %publish_packet, %c8 : index + %token_base = index.add %token_start, %token_offset : index + %token = index.add %token_base, %publish_token : index + %subgroup_channel_base = index.add %channel_base, %subgroup_channel_add : index + %channel = index.add %subgroup_channel_base, %publish_channel_add : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %writes = scalar.andi %valid_token, %valid_channel : i1 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + vector.fragment.store %result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes { + %bounded_token, %published_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%published_token_count]x[%bounded_output_size]xf32> + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %wide = vector.extf %values : vector<8xf16> to vector<8xf32> + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> + %published = scf.if %accumulates_output -> (vector<8xf32>) { + %residual = vector.load.mask %bounded_output_view[%bounded_token, %channel], %mask, %c0_f32x8 : view<[%published_token_count]x[%bounded_output_size]xf32>, vector<8xi1>, vector<8xf32> + %sum = vector.addf %residual, %wide : vector<8xf32> + scf.yield %sum : vector<8xf32> + } else { + scf.yield %wide : vector<8xf32> + } + vector.store.mask %published, %bounded_output_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%published_token_count]x[%bounded_output_size]xf32>, vector<8xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.store_single_row_overwrite_wave32> device @ggml_matmul_vector_q6_f32_store_single_row_overwrite_wave32(%result: vector<8xf32>, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %channel = index.add %channel_base, %subgroup_channel_add : index + %lane = kernel.subgroup.lane.id : index + %lane_token = index.rem %lane, %c16 : index + %lane_channel = index.div %lane, %c16 : index + %partial_channel_base = index.add %channel, %lane_channel : index + %stores_token0 = index.cmp eq, %lane_token, %c0 : index + scf.if %stores_token0 { + %dst = buffer.view %output[%base] : buffer -> view<1x[%output_size]xf32> + scf.for %element = [%c0 to %c8 step %c1] unroll { + %channel_step = index.mul %element, %c2 : index + %partial_channel = index.add %partial_channel_base, %channel_step : index + %channel_valid = index.cmp ult, %partial_channel, %output_size : index + scf.if %channel_valid { + %bounded_channel, %output_channels = index.assume %partial_channel, %output_size [lt(%partial_channel, %output_size)] : index, index + %value = vector.extract %result[%element] : vector<8xf32> -> f32 + view.store %value, %dst[%c0, %bounded_channel] : f32, view<1x[%output_size]xf32> + } + } + } + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.store_dense_outputs_wave32> device @ggml_matmul_vector_q6_f32_store_dense_outputs_wave32(%result: vector<8xf32>, %token_offset: index, %token_start: index, %token_count: index, %output_size: index, %channel_base: index, %subgroup_channel_add: index, %output: buffer, %output_accumulation: index) { + %base = index.constant 0 : offset + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %token = index.add %token_start, %token_offset : index + %channel = index.add %channel_base, %subgroup_channel_add : index + %token_end = index.add %token, %c16 : index + %channel_end = index.add %channel, %c16 : index + %valid_token = index.cmp ule, %token_end, %token_count : index + %valid_channel = index.cmp ule, %channel_end, %output_size : index + %valid = scalar.andi %valid_token, %valid_channel : i1 + scf.if %valid { + %layout = encoding.layout.strided [1, %output_size] : encoding + %dst = buffer.view %output[%base] : buffer -> view<[%output_size]x[%token_count]xf32, %layout> + %accumulates = index.cmp eq, %output_accumulation, %c1 : index + %published = scf.if %accumulates -> (vector<8xf32>) { + %old = vector.fragment.load %dst[%channel, %token] shape [%c16, %c16] : view<[%output_size]x[%token_count]xf32, %layout> -> vector<8xf32> + %sum = vector.addf %old, %result : vector<8xf32> + scf.yield %sum : vector<8xf32> + } else { + scf.yield %result : vector<8xf32> + } + vector.fragment.store %published, %dst[%channel, %token] shape [%c16, %c16] : vector<8xf32>, view<[%output_size]x[%token_count]xf32, %layout> + } else { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %lane = kernel.subgroup.lane.id : index + %lane_token = index.rem %lane, %c16 : index + %lane_channel = index.div %lane, %c16 : index + %partial_token = index.add %token, %lane_token : index + %partial_channel_base = index.add %channel, %lane_channel : index + %token_valid = index.cmp ult, %partial_token, %token_count : index + %accumulates = index.cmp eq, %output_accumulation, %c1 : index + scf.if %token_valid { + %bounded_token, %published_token_count = index.assume %partial_token, %token_count [lt(%partial_token, %token_count)] : index, index + %dst = buffer.view %output[%base] : buffer -> view<[%published_token_count]x[%output_size]xf32> + scf.for %element = [%c0 to %c8 step %c1] unroll { + %channel_step = index.mul %element, %c2 : index + %partial_channel = index.add %partial_channel_base, %channel_step : index + %channel_valid = index.cmp ult, %partial_channel, %output_size : index + scf.if %channel_valid { + %bounded_channel, %output_channels = index.assume %partial_channel, %output_size [lt(%partial_channel, %output_size)] : index, index + %value = vector.extract %result[%element] : vector<8xf32> -> f32 + %published = scf.if %accumulates -> (f32) { + %old = view.load %dst[%bounded_token, %bounded_channel] : view<[%published_token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %value, %old : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + view.store %published, %dst[%bounded_token, %bounded_channel] : f32, view<[%published_token_count]x[%output_size]xf32> + } + } + } + } + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.body_packed_storage> device @ggml_matmul_vector_q6_f32_body_packed_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c13440_bytes = index.constant 13440 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %empty_words = vector.constant 0 : vector<4xi32> + %zero_scale_bytes = vector.constant 0 : vector<8xi8> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high0_mask = vector.constant 50529027 : vector<4xi32> + %high1_mask = vector.constant 202116108 : vector<4xi32> + %high2_mask = vector.constant 808464432 : vector<4xi32> + %high3_mask = vector.constant -1061109568 : vector<4xi32> + %shift2 = vector.constant 2 : vector<4xi32> + %shift4 = vector.constant 4 : vector<4xi32> + %negative32_f32 = vector.constant -32.0 : vector<16xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x16xf16, %activation_fragment_layout> + + %channel_base = index.mul %output_block, %c64 : index + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %channel = index.add %channel_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %use_scale_row_layout = index.cmp eq, %scale_row_layout, %c1 : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<8xf32>) -> (vector<8xf32>) { + %row_group_block0 = index.mul %output_block, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %row_group_byte_base = index.scale %row_group_block, %c13440_bytes : index, offset -> offset + %scale_byte_base = index.add %row_group_byte_base, %c128_bytes : offset + %raw_byte_base = index.add %row_group_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%row_group_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %scale_row_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_noalias[%raw_byte_base] : buffer -> view<2x3x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %load_scale_row = scalar.andi %valid_channel, %use_scale_row_layout : i1 + %scale_bytes = scf.if %load_scale_row -> (vector<8xi8>) { + %loaded_scale_bytes = vector.load %scale_row_view[%weight_local_row, %weight_half, %c0] : view<64x2x8xi8> -> vector<8xi8> + scf.yield %loaded_scale_bytes : vector<8xi8> + } else { + scf.yield %zero_scale_bytes : vector<8xi8> + } + %block_result, %final_ql0_words, %final_ql1_words, %final_qh_words = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<8xf32>, %cached_ql0_words = %empty_words : vector<4xi32>, %cached_ql1_words = %empty_words : vector<4xi32>, %cached_qh_words = %empty_words : vector<4xi32>) -> (vector<8xf32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %quant_group_next = index.add %quant_group, %c1 : index + %half128 = index.div %quant_group, %c4 : index + %group_in_half = index.rem %quant_group, %c4 : index + %is_low_pair = index.cmp eq, %group_in_half, %c0 : index + %is_first_group = index.cmp eq, %quant_group, %c0 : index + %prefetch_second_half = index.cmp eq, %quant_group, %c2 : index + %weight_values0, %weight_values1, %next_ql0_words, %next_ql1_words, %next_qh_words = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_words, %ql1_words, %qh_words = scf.if %is_first_group -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_bytes = vector.load %raw_view[%half128, %c0, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %ql1_bytes = vector.load %raw_view[%half128, %c1, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %qh_bytes = vector.load %raw_view[%half128, %c2, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %loaded_ql0_words = vector.bitcast %ql0_bytes : vector<16xi8> to vector<4xi32> + %loaded_ql1_words = vector.bitcast %ql1_bytes : vector<16xi8> to vector<4xi32> + %loaded_qh_words = vector.bitcast %qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %loaded_ql0_words, %loaded_ql1_words, %loaded_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %codes0_words, %codes1_words = scf.if %is_low_pair -> (vector<4xi32>, vector<4xi32>) { + %ql0_low = vector.andi %ql0_words, %nibble_mask : vector<4xi32> + %ql1_low = vector.andi %ql1_words, %nibble_mask : vector<4xi32> + %qh0_bits0 = vector.andi %qh_words, %high0_mask : vector<4xi32> + %qh1_bits0 = vector.andi %qh_words, %high1_mask : vector<4xi32> + %qh0_bits = vector.shli %qh0_bits0, %shift4 : vector<4xi32> + %qh1_bits = vector.shli %qh1_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_low, %qh0_bits : vector<4xi32> + %codes1 = vector.ori %ql1_low, %qh1_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } else { + %ql0_high0 = vector.shrui %ql0_words, %shift4 : vector<4xi32> + %ql1_high0 = vector.shrui %ql1_words, %shift4 : vector<4xi32> + %ql0_high = vector.andi %ql0_high0, %nibble_mask : vector<4xi32> + %ql1_high = vector.andi %ql1_high0, %nibble_mask : vector<4xi32> + %qh2_bits = vector.andi %qh_words, %high2_mask : vector<4xi32> + %qh3_bits0 = vector.andi %qh_words, %high3_mask : vector<4xi32> + %qh3_bits = vector.shrui %qh3_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_high, %qh2_bits : vector<4xi32> + %codes1 = vector.ori %ql1_high, %qh3_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } + %codes0_u8 = vector.bitcast %codes0_words : vector<4xi32> to vector<16xi8> + %codes1_u8 = vector.bitcast %codes1_words : vector<4xi32> to vector<16xi8> + %codes0_unsigned_f32 = vector.uitofp %codes0_u8 : vector<16xi8> to vector<16xf32> + %codes1_unsigned_f32 = vector.uitofp %codes1_u8 : vector<16xi8> to vector<16xf32> + %codes0_f32 = vector.addf %codes0_unsigned_f32, %negative32_f32 : vector<16xf32> + %codes1_f32 = vector.addf %codes1_unsigned_f32, %negative32_f32 : vector<16xf32> + %scale0_i8, %scale1_i8 = scf.if %use_scale_row_layout -> (i8, i8) { + %row_scale0_i8 = vector.extract %scale_bytes[%quant_group] : vector<8xi8> -> i8 + %row_scale1_i8 = vector.extract %scale_bytes[%quant_group_next] : vector<8xi8> -> i8 + scf.yield %row_scale0_i8, %row_scale1_i8 : i8, i8 + } else { + %group_scale0_i8 = view.load %scale_view[%quant_group, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %group_scale1_i8 = view.load %scale_view[%quant_group_next, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + scf.yield %group_scale0_i8, %group_scale1_i8 : i8, i8 + } + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words = scf.if %prefetch_second_half -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %next_ql0_bytes = vector.load %raw_view[%c1, %c0, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql1_bytes = vector.load %raw_view[%c1, %c1, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_qh_bytes = vector.load %raw_view[%c1, %c2, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql0 = vector.bitcast %next_ql0_bytes : vector<16xi8> to vector<4xi32> + %next_ql1 = vector.bitcast %next_ql1_bytes : vector<16xi8> to vector<4xi32> + %next_qh = vector.bitcast %next_qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %next_ql0, %next_ql1, %next_qh : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %ql0_words, %ql1_words, %qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %values0, %values1, %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1, %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<8xf32>) -> (vector<8xf32>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs, %rhs, %half_acc : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.yield %half_next : vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next, %next_ql0_words, %next_ql1_words, %next_qh_words : vector<8xf32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %block_result : vector<8xf32> + } + + %overwrites_output = index.cmp eq, %output_accumulation, %c0 : index + %has_one_token = index.cmp eq, %bounded_token_count, %c1 : index + %use_single_token_overwrite = scalar.andi %overwrites_output, %has_one_token : i1 + scf.if %use_single_token_overwrite { + template.apply<@ggml.matmul.vector_q6_f32.store_single_row_overwrite_wave32>(%result, %bounded_output_size, %channel_base, %subgroup_channel_add, %output_noalias) : (vector<8xf32>, index, index, index, buffer) + } else { + template.apply<@ggml.matmul.vector_q6_f32.store_dense_outputs_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + } + template.return +} + +template.def<@ggml.matmul.vector_q6_f32.body_i8_storage> device @ggml_matmul_vector_q6_f32_body_i8_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c17536_bytes = index.constant 17536 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x16xf16, %activation_fragment_layout> + + %channel_base = index.mul %output_block, %c64 : index + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %channel = index.add %channel_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<8xf32>) -> (vector<8xf32>) { + %row_group_block0 = index.mul %output_block, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %row_group_byte_base = index.scale %row_group_block, %c17536_bytes : index, offset -> offset + %scale_byte_base = index.add %row_group_byte_base, %c128_bytes : offset + %code_byte_base = index.add %row_group_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%row_group_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %code_view = buffer.view %weight_noalias[%code_byte_base] : buffer -> view<8x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %block_result = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<8xf32>) -> (vector<8xf32>) { + %quant_group_next = index.add %quant_group, %c1 : index + %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { + %codes0 = vector.load %code_view[%quant_group, %weight_local_row, %weight_k] : view<8x64x32xi8> -> vector<16xi8> + %codes1 = vector.load %code_view[%quant_group_next, %weight_local_row, %weight_k] : view<8x64x32xi8> -> vector<16xi8> + %codes0_f32 = vector.sitofp %codes0 : vector<16xi8> to vector<16xf32> + %codes1_f32 = vector.sitofp %codes1 : vector<16xi8> to vector<16xf32> + %scale0_i8 = view.load %scale_view[%quant_group, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %scale1_i8 = view.load %scale_view[%quant_group_next, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + scf.yield %values0, %values1 : vector<16xf16>, vector<16xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1 : vector<16xf16>, vector<16xf16> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<8xf32>) -> (vector<8xf32>) unroll { + %lhs_fragment = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs_fragment = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs_fragment, %rhs_fragment, %half_acc : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.yield %half_next : vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next : vector<8xf32> + } + scf.yield %block_result : vector<8xf32> + } + + template.apply<@ggml.matmul.vector_q6_f32.store_dense_outputs_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.return +} + + +template.def<@ggml.matmul.vector_q6_f32.body_selected_outputs> device @ggml_matmul_vector_q6_f32_body_selected_outputs(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %selected_count0: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %weight_offset_index = config.get @ggml.matmul.vector.weight_offset : index + %weight_offset = index.cast %weight_offset_index : index to offset + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %bounded_selected_count = index.assume %selected_count0 [range(%selected_count0, 1, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c13440_bytes = index.constant 13440 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %result_stage_bytes = index.constant 2048 : offset + %zero_accumulator = vector.constant 0.0 : vector<16xf16> + %empty_words = vector.constant 0 : vector<4xi32> + %zero_scale_bytes = vector.constant 0 : vector<8xi8> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high0_mask = vector.constant 50529027 : vector<4xi32> + %high1_mask = vector.constant 202116108 : vector<4xi32> + %high2_mask = vector.constant 808464432 : vector<4xi32> + %high3_mask = vector.constant -1061109568 : vector<4xi32> + %shift2 = vector.constant 2 : vector<4xi32> + %shift4 = vector.constant 4 : vector<4xi32> + %negative32_f32 = vector.constant -32.0 : vector<16xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %selected_outputs_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %selected_outputs, %output : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%base] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %selected_outputs_view = buffer.view %selected_outputs_noalias[%base] : buffer -> view<64xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%base] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%base] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%base] : buffer -> view<64x16xf16, %activation_fragment_layout> + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %valid_selected_output = index.cmp ult, %weight_local_row, %bounded_selected_count : index + %channel0 = scf.if %valid_selected_output -> (index) { + %selected_output_i32 = view.load %selected_outputs_view[%weight_local_row] : view<64xi32> -> i32 + %selected_output_index0 = index.cast %selected_output_i32 : i32 to index + %selected_output_index = index.assume %selected_output_index0 [range(%selected_output_index0, 0, 262143)] : index + scf.yield %selected_output_index : index + } else { + scf.yield %c0 : index + } + %channel = index.assume %channel0 [range(%channel0, 0, 262143)] : index + %channel_in_range = index.cmp ult, %channel, %bounded_output_size : index + %valid_channel = scalar.andi %valid_selected_output, %channel_in_range : i1 + %selected_output_block = index.div %channel, %c64 : index + %selected_weight_local_row = index.rem %channel, %c64 : index + %use_scale_row_layout = index.cmp eq, %scale_row_layout, %c1 : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<16xf16>) -> (vector<16xf16>) { + %row_group_block0 = index.mul %selected_output_block, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %row_group_byte_relative = index.scale %row_group_block, %c13440_bytes : index, offset -> offset + %row_group_byte_base = index.add %weight_offset, %row_group_byte_relative : offset + %scale_byte_base = index.add %row_group_byte_base, %c128_bytes : offset + %raw_byte_base = index.add %row_group_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%row_group_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %scale_row_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_noalias[%raw_byte_base] : buffer -> view<2x3x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%selected_weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %zero_f32 = scalar.constant 0.0 : f32 + scf.yield %zero_f32 : f32 + } + %load_scale_row = scalar.andi %valid_channel, %use_scale_row_layout : i1 + %scale_bytes = scf.if %load_scale_row -> (vector<8xi8>) { + %loaded_scale_bytes = vector.load %scale_row_view[%selected_weight_local_row, %weight_half, %c0] : view<64x2x8xi8> -> vector<8xi8> + scf.yield %loaded_scale_bytes : vector<8xi8> + } else { + scf.yield %zero_scale_bytes : vector<8xi8> + } + %block_result, %final_ql0_words, %final_ql1_words, %final_qh_words = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<16xf16>, %cached_ql0_words = %empty_words : vector<4xi32>, %cached_ql1_words = %empty_words : vector<4xi32>, %cached_qh_words = %empty_words : vector<4xi32>) -> (vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %quant_group_next = index.add %quant_group, %c1 : index + %half128 = index.div %quant_group, %c4 : index + %group_in_half = index.rem %quant_group, %c4 : index + %is_low_pair = index.cmp eq, %group_in_half, %c0 : index + %is_first_group = index.cmp eq, %quant_group, %c0 : index + %prefetch_second_half = index.cmp eq, %quant_group, %c2 : index + %weight_values0, %weight_values1, %next_ql0_words, %next_ql1_words, %next_qh_words = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_words, %ql1_words, %qh_words = scf.if %is_first_group -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_bytes = vector.load %raw_view[%half128, %c0, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %ql1_bytes = vector.load %raw_view[%half128, %c1, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %qh_bytes = vector.load %raw_view[%half128, %c2, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %loaded_ql0_words = vector.bitcast %ql0_bytes : vector<16xi8> to vector<4xi32> + %loaded_ql1_words = vector.bitcast %ql1_bytes : vector<16xi8> to vector<4xi32> + %loaded_qh_words = vector.bitcast %qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %loaded_ql0_words, %loaded_ql1_words, %loaded_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %codes0_words, %codes1_words = scf.if %is_low_pair -> (vector<4xi32>, vector<4xi32>) { + %ql0_low = vector.andi %ql0_words, %nibble_mask : vector<4xi32> + %ql1_low = vector.andi %ql1_words, %nibble_mask : vector<4xi32> + %qh0_bits0 = vector.andi %qh_words, %high0_mask : vector<4xi32> + %qh1_bits0 = vector.andi %qh_words, %high1_mask : vector<4xi32> + %qh0_bits = vector.shli %qh0_bits0, %shift4 : vector<4xi32> + %qh1_bits = vector.shli %qh1_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_low, %qh0_bits : vector<4xi32> + %codes1 = vector.ori %ql1_low, %qh1_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } else { + %ql0_high0 = vector.shrui %ql0_words, %shift4 : vector<4xi32> + %ql1_high0 = vector.shrui %ql1_words, %shift4 : vector<4xi32> + %ql0_high = vector.andi %ql0_high0, %nibble_mask : vector<4xi32> + %ql1_high = vector.andi %ql1_high0, %nibble_mask : vector<4xi32> + %qh2_bits = vector.andi %qh_words, %high2_mask : vector<4xi32> + %qh3_bits0 = vector.andi %qh_words, %high3_mask : vector<4xi32> + %qh3_bits = vector.shrui %qh3_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_high, %qh2_bits : vector<4xi32> + %codes1 = vector.ori %ql1_high, %qh3_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } + %codes0_u8 = vector.bitcast %codes0_words : vector<4xi32> to vector<16xi8> + %codes1_u8 = vector.bitcast %codes1_words : vector<4xi32> to vector<16xi8> + %codes0_unsigned_f32 = vector.uitofp %codes0_u8 : vector<16xi8> to vector<16xf32> + %codes1_unsigned_f32 = vector.uitofp %codes1_u8 : vector<16xi8> to vector<16xf32> + %codes0_f32 = vector.addf %codes0_unsigned_f32, %negative32_f32 : vector<16xf32> + %codes1_f32 = vector.addf %codes1_unsigned_f32, %negative32_f32 : vector<16xf32> + %scale0_i8, %scale1_i8 = scf.if %use_scale_row_layout -> (i8, i8) { + %row_scale0_i8 = vector.extract %scale_bytes[%quant_group] : vector<8xi8> -> i8 + %row_scale1_i8 = vector.extract %scale_bytes[%quant_group_next] : vector<8xi8> -> i8 + scf.yield %row_scale0_i8, %row_scale1_i8 : i8, i8 + } else { + %group_scale0_i8 = view.load %scale_view[%quant_group, %selected_weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %group_scale1_i8 = view.load %scale_view[%quant_group_next, %selected_weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + scf.yield %group_scale0_i8, %group_scale1_i8 : i8, i8 + } + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words = scf.if %prefetch_second_half -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %next_ql0_bytes = vector.load %raw_view[%c1, %c0, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql1_bytes = vector.load %raw_view[%c1, %c1, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_qh_bytes = vector.load %raw_view[%c1, %c2, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql0 = vector.bitcast %next_ql0_bytes : vector<16xi8> to vector<4xi32> + %next_ql1 = vector.bitcast %next_ql1_bytes : vector<16xi8> to vector<4xi32> + %next_qh = vector.bitcast %next_qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %next_ql0, %next_ql1, %next_qh : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %ql0_words, %ql1_words, %qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %values0, %values1, %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1, %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs, %rhs, %half_acc : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %half_next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next, %next_ql0_words, %next_ql1_words, %next_qh_words : vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %block_result : vector<16xf16> + } + template.apply<@ggml.matmul.vector_q6_f32.store_fragment_outputs_wave32>(%result, %c0, %c0, %bounded_token_count, %c64, %c0, %subgroup_channel_add, %result_stage, %output_noalias, %output_accumulation) : (vector<16xf16>, index, index, index, index, index, index, buffer, buffer, index) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_down_schedule.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_down_schedule.loom new file mode 100644 index 000000000000..a49ab9e5c4ce --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_down_schedule.loom @@ -0,0 +1,375 @@ +// Expert-grouped Q4_K and Q6_K routed down projections for gfx11 prefill. +// +// Each two-wave workgroup decodes 64 output channels for 32 routed rows of one +// expert. Raw GGUF weights and FP16 routed SwiGLU rows are staged as FP16 WMMA +// operands. Results remain FP16 and are scattered into compact +// [token, route, hidden] order without collisions. A following reduction +// widens each route, applies its normalized weight, and accumulates the +// residual in FP32. +// +// Top-k routing selects an expert at most once per token. The expert table +// therefore contains at most token_count assignments per expert, while each +// assignment ordinal still identifies the compact [token, route] activation +// and route-weight row. +template.decl @ggml.mul_mat_id_f16_f16_wmma.core(%weight_format: index, %token_count: index, %input_size: index, %route_count: index, %expert_count: index, %output_size: index, %channel_tile: index, %route_tile: index, %expert: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_id_f16_f16_wmma.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %channel_tile: index, %expert: index, %route_tile_base: index, %expert_route_count: index, %expert_table: buffer, %output: buffer) + +template.decl @ggml.mul_mat_id_f16_f16_wmma.publish_vector4(%publish_word: i1, %assignment: index, %channel: index, %token_count: index, %route_count: index, %output_size: index, %values: vector<4xf16>, %output: buffer) + +func.decl @ggml_q4k_split_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) + +func.decl @ggml_q4k_split_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) + +func.decl @ggml_q4k_split_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) + +func.decl @ggml_q6k_split_wmma_load_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_half: index, %packet: index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) + +func.decl @ggml_q6k_split_wmma_vector4_from_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index, %ql0_word: vector<1xi32>, %ql1_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) + +// Shared raw-quantized matrix schedule. Entry points pass a literal weight +// format so linking and JIT specialization erase the inactive packed decoder. +template.def<@ggml.mul_mat_id_f16_f16_wmma.core> device @ggml_mul_mat_id_f16_f16_wmma_core(%weight_format: index, %token_count: index, %input_size: index, %route_count: index, %expert_count: index, %output_size: index, %channel_tile: index, %route_tile: index, %expert: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) { + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c6 = index.constant 6 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : 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 + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %q4_block_bytes = index.constant 144 : offset + %q6_block_bytes = index.constant 210 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_i32x1 = vector.constant 0 : vector<1xi32> + %c0_i32x4 = vector.constant 0 : vector<4xi32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %token_count, %bounded_route_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q6 = index.cmp eq, %weight_format, %c6 : index + %quant_block_bytes = scf.select %is_q6, %q6_block_bytes, %q4_block_bytes : offset + %quant_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %quant_block_count, %quant_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %weight_noalias = buffer.assume.noalias %input, %expert_table, %weight : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%input_size]xf16> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %initial_route_tile_base = index.mul %route_tile, %c32 : index + // The launch workload fixes the interleaved route partitions for this exact + // command-program specialization. + %padded_token_count = index.add %token_count, %c63 : index + %route_partition_count = index.div %padded_token_count, %c64 : index + %route_partition_step = index.mul %route_partition_count, %c32 : index + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 + %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 + %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index + %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index + scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { + // Snapshot this expert's compact assignment map once, then reuse it for + // every K group and output-channel tile. + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + // Four lanes cover the 64 output rows owned by one load packet. Snapshot + // each Q4_K block header once and retain it across all eight quant groups. + // The Q6 specialization erases this loop with %is_q4=false. + %header0, %header1, %header2, %header3 = scf.for %header_row_offset = [%c0 to %c64 step %c16](%prior_header0 = %c0_i32x4 : vector<4xi32>, %prior_header1 = %c0_i32x4 : vector<4xi32>, %prior_header2 = %c0_i32x4 : vector<4xi32>, %prior_header3 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %header_local_row0 = index.add %load_row, %header_row_offset : index + %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index + %header_channel = index.add %channel_tile_base, %header_local_row : index + %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index + %loads_q4_header = scalar.andi %is_q4, %valid_header_channel : i1 + %loaded_header = scf.if %loads_q4_header -> (vector<4xi32>) { + %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset + %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset + %header_words = func.call @ggml_q4k_split_wmma_load_header(%weight_noalias, %header_row_byte_base, %quant_block) : (buffer, offset, index) -> (vector<4xi32>) + scf.yield %header_words : vector<4xi32> + } else { + scf.yield %c0_i32x4 : vector<4xi32> + } + %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index + %updates_header1 = index.cmp eq, %header_row_offset, %c16 : index + %updates_header2 = index.cmp eq, %header_row_offset, %c32 : index + %updates_header3 = index.cmp eq, %header_row_offset, %c48 : index + %next_header0 = scf.select %updates_header0, %loaded_header, %prior_header0 : vector<4xi32> + %next_header1 = scf.select %updates_header1, %loaded_header, %prior_header1 : vector<4xi32> + %next_header2 = scf.select %updates_header2, %loaded_header, %prior_header2 : vector<4xi32> + %next_header3 = scf.select %updates_header3, %loaded_header, %prior_header3 : vector<4xi32> + scf.yield %next_header0, %next_header1, %next_header2, %next_header3 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %quant_group_outer_count = scf.select %is_q4, %c4, %c2 : index + %groups_per_outer = scf.select %is_q4, %c2, %c4 : index + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group_outer = [%c0 to %quant_group_outer_count 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>) { + // Q4_K retains one packed code word across its adjacent low/high group + // pair. The Q6 specialization erases this acquisition path. + %q_word0, %q_word1, %q_word2, %q_word3 = scf.for %code_row_offset = [%c0 to %c64 step %c16](%prior_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_q_word2 = %c0_i32x1 : vector<1xi32>, %prior_q_word3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %code_local_row0 = index.add %load_row, %code_row_offset : index + %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index + %code_channel = index.add %channel_tile_base, %code_local_row : index + %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index + %loads_q4_code = scalar.andi %is_q4, %valid_code_channel : i1 + %loaded_q_word = scf.if %loads_q4_code -> (vector<1xi32>) { + %q4_group_pair = index.assume %quant_group_outer [range(%quant_group_outer, 0, 3)] : index + %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset + %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset + %q_word = func.call @ggml_q4k_split_wmma_load_code(%weight_noalias, %code_row_byte_base, %quant_block, %q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + scf.yield %q_word : vector<1xi32> + } else { + scf.yield %c0_i32x1 : vector<1xi32> + } + %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index + %updates_q_word1 = index.cmp eq, %code_row_offset, %c16 : index + %updates_q_word2 = index.cmp eq, %code_row_offset, %c32 : index + %updates_q_word3 = index.cmp eq, %code_row_offset, %c48 : index + %next_q_word0 = scf.select %updates_q_word0, %loaded_q_word, %prior_q_word0 : vector<1xi32> + %next_q_word1 = scf.select %updates_q_word1, %loaded_q_word, %prior_q_word1 : vector<1xi32> + %next_q_word2 = scf.select %updates_q_word2, %loaded_q_word, %prior_q_word2 : vector<1xi32> + %next_q_word3 = scf.select %updates_q_word3, %loaded_q_word, %prior_q_word3 : vector<1xi32> + scf.yield %next_q_word0, %next_q_word1, %next_q_word2, %next_q_word3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + // Four Q6_K groups in one half-block share a QH word and consume two + // nibbles from each of two QL words. Retain the three packed words for + // all four groups. The Q4 specialization erases this acquisition path. + %q6_ql00, %q6_ql10, %q6_qh0, %q6_ql01, %q6_ql11, %q6_qh1, %q6_ql02, %q6_ql12, %q6_qh2, %q6_ql03, %q6_ql13, %q6_qh3 = scf.for %q6_code_row_offset = [%c0 to %c64 step %c16](%prior_q6_ql00 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql10 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh0 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql01 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql11 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh1 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql02 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql12 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh2 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql03 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql13 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %q6_code_local_row0 = index.add %load_row, %q6_code_row_offset : index + %q6_code_local_row = index.assume %q6_code_local_row0 [range(%q6_code_local_row0, 0, 63)] : index + %q6_code_channel = index.add %channel_tile_base, %q6_code_local_row : index + %valid_q6_code_channel = index.cmp ult, %q6_code_channel, %bounded_output_size : index + %loads_q6_code = scalar.andi %is_q6, %valid_q6_code_channel : i1 + %loaded_q6_ql0, %loaded_q6_ql1, %loaded_q6_qh = scf.if %loads_q6_code -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { + %q6_half = index.assume %quant_group_outer [range(%quant_group_outer, 0, 1)] : index + %q6_code_channel_byte_add = index.scale %q6_code_channel, %weight_row_bytes : index, offset -> offset + %q6_code_row_byte_base = index.add %expert_byte_base, %q6_code_channel_byte_add : offset + %ql0_word, %ql1_word, %qh_word = func.call @ggml_q6k_split_wmma_load_half_codes(%weight_noalias, %q6_code_row_byte_base, %quant_block, %q6_half, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) + scf.yield %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> + } else { + scf.yield %c0_i32x1, %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %updates_q6_code0 = index.cmp eq, %q6_code_row_offset, %c0 : index + %updates_q6_code1 = index.cmp eq, %q6_code_row_offset, %c16 : index + %updates_q6_code2 = index.cmp eq, %q6_code_row_offset, %c32 : index + %updates_q6_code3 = index.cmp eq, %q6_code_row_offset, %c48 : index + %next_q6_ql00 = scf.select %updates_q6_code0, %loaded_q6_ql0, %prior_q6_ql00 : vector<1xi32> + %next_q6_ql10 = scf.select %updates_q6_code0, %loaded_q6_ql1, %prior_q6_ql10 : vector<1xi32> + %next_q6_qh0 = scf.select %updates_q6_code0, %loaded_q6_qh, %prior_q6_qh0 : vector<1xi32> + %next_q6_ql01 = scf.select %updates_q6_code1, %loaded_q6_ql0, %prior_q6_ql01 : vector<1xi32> + %next_q6_ql11 = scf.select %updates_q6_code1, %loaded_q6_ql1, %prior_q6_ql11 : vector<1xi32> + %next_q6_qh1 = scf.select %updates_q6_code1, %loaded_q6_qh, %prior_q6_qh1 : vector<1xi32> + %next_q6_ql02 = scf.select %updates_q6_code2, %loaded_q6_ql0, %prior_q6_ql02 : vector<1xi32> + %next_q6_ql12 = scf.select %updates_q6_code2, %loaded_q6_ql1, %prior_q6_ql12 : vector<1xi32> + %next_q6_qh2 = scf.select %updates_q6_code2, %loaded_q6_qh, %prior_q6_qh2 : vector<1xi32> + %next_q6_ql03 = scf.select %updates_q6_code3, %loaded_q6_ql0, %prior_q6_ql03 : vector<1xi32> + %next_q6_ql13 = scf.select %updates_q6_code3, %loaded_q6_ql1, %prior_q6_ql13 : vector<1xi32> + %next_q6_qh3 = scf.select %updates_q6_code3, %loaded_q6_qh, %prior_q6_qh3 : vector<1xi32> + scf.yield %next_q6_ql00, %next_q6_ql10, %next_q6_qh0, %next_q6_ql01, %next_q6_ql11, %next_q6_qh1, %next_q6_ql02, %next_q6_ql12, %next_q6_qh2, %next_q6_ql03, %next_q6_ql13, %next_q6_qh3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %outer_result00, %outer_result01, %outer_result10, %outer_result11 = scf.for %group_within_outer = [%c0 to %groups_per_outer step %c1](%group_acc00 = %acc00 : vector<8xf16>, %group_acc01 = %acc01 : vector<8xf16>, %group_acc10 = %acc10 : vector<8xf16>, %group_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %quant_group_base = index.mul %quant_group_outer, %groups_per_outer : index + %quant_group0 = index.add %quant_group_base, %group_within_outer : index + %quant_group = index.assume %quant_group0 [range(%quant_group0, 0, 7)] : index + %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 + %selects_row0 = index.cmp eq, %row_offset, %c0 : index + %selects_row2 = index.cmp eq, %row_offset, %c32 : index + %selects_low_row_pair = index.cmp ult, %row_offset, %c32 : index + %selected_header01 = scf.select %selects_row0, %header0, %header1 : vector<4xi32> + %selected_header23 = scf.select %selects_row2, %header2, %header3 : vector<4xi32> + %selected_header = scf.select %selects_low_row_pair, %selected_header01, %selected_header23 : vector<4xi32> + %selected_q_word01 = scf.select %selects_row0, %q_word0, %q_word1 : vector<1xi32> + %selected_q_word23 = scf.select %selects_row2, %q_word2, %q_word3 : vector<1xi32> + %selected_q_word = scf.select %selects_low_row_pair, %selected_q_word01, %selected_q_word23 : vector<1xi32> + %selected_q6_ql0_01 = scf.select %selects_row0, %q6_ql00, %q6_ql01 : vector<1xi32> + %selected_q6_ql0_23 = scf.select %selects_row2, %q6_ql02, %q6_ql03 : vector<1xi32> + %selected_q6_ql0 = scf.select %selects_low_row_pair, %selected_q6_ql0_01, %selected_q6_ql0_23 : vector<1xi32> + %selected_q6_ql1_01 = scf.select %selects_row0, %q6_ql10, %q6_ql11 : vector<1xi32> + %selected_q6_ql1_23 = scf.select %selects_row2, %q6_ql12, %q6_ql13 : vector<1xi32> + %selected_q6_ql1 = scf.select %selects_low_row_pair, %selected_q6_ql1_01, %selected_q6_ql1_23 : vector<1xi32> + %selected_q6_qh01 = scf.select %selects_row0, %q6_qh0, %q6_qh1 : vector<1xi32> + %selected_q6_qh23 = scf.select %selects_row2, %q6_qh2, %q6_qh3 : vector<1xi32> + %selected_q6_qh = scf.select %selects_low_row_pair, %selected_q6_qh01, %selected_q6_qh23 : vector<1xi32> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %decoded = scf.if %is_q6 -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %q6_values = func.call @ggml_q6k_split_wmma_vector4_from_half_codes(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet, %selected_q6_ql0, %selected_q6_ql1, %selected_q6_qh) : (buffer, offset, index, index, index, vector<1xi32>, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) + scf.yield %q6_values : vector<4xf16> + } else { + %q4_values = func.call @ggml_q4k_split_wmma_vector4_from_header_code(%quant_group, %selected_header, %selected_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) + scf.yield %q4_values : vector<4xf16> + } + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%bounded_assignment, %input_k] : view<[%assignment_count]x[%input_size]xf16> -> vector<4xf16> + scf.yield %loaded : 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 = %group_acc00 : vector<8xf16>, %half_acc01 = %group_acc01 : vector<8xf16>, %half_acc10 = %group_acc10 : vector<8xf16>, %half_acc11 = %group_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 %outer_result00, %outer_result01, %outer_result10, %outer_result11 : 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> + } + // WMMA produces [channel][route] fragments. Transpose through wave-private + // LDS so each lane publishes four adjacent channels for one assignment. + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.select %assignment0_nonnegative, %assignment0_i32, %c0_i32 : i32 + %safe_assignment1_i32 = scf.select %assignment1_nonnegative, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + %values00 = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + template.apply<@ggml.mul_mat_id_f16_f16_wmma.publish_vector4>(%writes00, %bounded_assignment0, %channel0, %token_count, %bounded_route_count, %bounded_output_size, %values00, %output) : (i1, index, index, index, index, index, vector<4xf16>, 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) + %values01 = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + template.apply<@ggml.mul_mat_id_f16_f16_wmma.publish_vector4>(%writes01, %bounded_assignment1, %channel0, %token_count, %bounded_route_count, %bounded_output_size, %values01, %output) : (i1, index, index, index, index, index, vector<4xf16>, 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) + %values10 = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + template.apply<@ggml.mul_mat_id_f16_f16_wmma.publish_vector4>(%writes10, %bounded_assignment0, %channel1, %token_count, %bounded_route_count, %bounded_output_size, %values10, %output) : (i1, index, index, index, index, index, vector<4xf16>, 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) + %values11 = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + template.apply<@ggml.mul_mat_id_f16_f16_wmma.publish_vector4>(%writes11, %bounded_assignment1, %channel1, %token_count, %bounded_route_count, %bounded_output_size, %values11, %output) : (i1, index, index, index, index, index, vector<4xf16>, buffer) + template.apply<@ggml.mul_mat_id_f16_f16_wmma.finish_tile>(%token_count, %bounded_route_count, %bounded_output_size, %bounded_expert_count, %channel_tile, %bounded_expert, %route_tile_base, %expert_route_count, %expert_table, %output) : (index, index, index, index, index, index, index, index, buffer, buffer) + // Route, operand, and result stages are reused by the next concentrated + // routing partition. All waves must finish publication before reuse. + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_accumulate.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_accumulate.loom new file mode 100644 index 000000000000..d656345c7f0a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_accumulate.loom @@ -0,0 +1,395 @@ +// Compile-time selected routed F16 pair accumulator. + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.accumulate_block(%quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %gate_weight_row_bytes: offset, %up_weight_row_bytes: offset, %gate_weight_format: index, %up_weight_format: index, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.accumulate_q4k_block(%quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.accumulate_q5k_block(%quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.accumulate_q6k_block(%quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.accumulate_generic_block(%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>, %quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %gate_weight_row_bytes: offset, %up_weight_row_bytes: offset, %gate_weight_format: index, %up_weight_format: index, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + +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>) + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_block> device @ggml_mul_mat_id_swiglu_f16_f16_accumulate_block(%quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %gate_weight_row_bytes: offset, %up_weight_row_bytes: offset, %gate_weight_format: index, %up_weight_format: index, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %c4 = index.constant 4 : index + %c5 = index.constant 5 : index + %c6 = index.constant 6 : index + %same_weight_format = index.cmp eq, %gate_weight_format, %up_weight_format : index + %gate_is_q4k = index.cmp eq, %gate_weight_format, %c4 : index + %gate_is_q5k = index.cmp eq, %gate_weight_format, %c5 : index + %gate_is_q6k = index.cmp eq, %gate_weight_format, %c6 : index + %is_q4k = scalar.andi %same_weight_format, %gate_is_q4k : i1 + %is_q5k = scalar.andi %same_weight_format, %gate_is_q5k : i1 + %is_q6k = scalar.andi %same_weight_format, %gate_is_q6k : i1 + %q4_gate, %q4_up = scf.if %is_q4k -> (vector<16xf16>, vector<16xf16>) { + %q4_gate, %q4_up = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_q4k_block>(%quant_block, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_weight, %up_weight, %expert_byte_base, %gate_weight_row_bytes, %gate_accumulator, %up_accumulator) : (index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<16xf16>, vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + scf.yield %q4_gate, %q4_up : vector<16xf16>, vector<16xf16> + } else { + scf.yield %gate_accumulator, %up_accumulator : vector<16xf16>, vector<16xf16> + } + %q5_gate, %q5_up = scf.if %is_q5k -> (vector<16xf16>, vector<16xf16>) { + %q5_gate, %q5_up = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_q5k_block>(%quant_block, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_weight, %up_weight, %expert_byte_base, %gate_weight_row_bytes, %gate_accumulator, %up_accumulator) : (index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<16xf16>, vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + scf.yield %q5_gate, %q5_up : vector<16xf16>, vector<16xf16> + } else { + scf.yield %gate_accumulator, %up_accumulator : vector<16xf16>, vector<16xf16> + } + %q6_gate, %q6_up = scf.if %is_q6k -> (vector<16xf16>, vector<16xf16>) { + %q6_gate, %q6_up = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_q6k_block>(%quant_block, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_weight, %up_weight, %expert_byte_base, %gate_weight_row_bytes, %gate_accumulator, %up_accumulator) : (index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<16xf16>, vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + scf.yield %q6_gate, %q6_up : vector<16xf16>, vector<16xf16> + } else { + scf.yield %gate_accumulator, %up_accumulator : vector<16xf16>, vector<16xf16> + } + %is_q4_or_q5 = scalar.ori %is_q4k, %is_q5k : i1 + %uses_fast_path = scalar.ori %is_q4_or_q5, %is_q6k : i1 + %generic_gate, %generic_up = scf.if %uses_fast_path -> (vector<16xf16>, vector<16xf16>) { + scf.yield %gate_accumulator, %up_accumulator : vector<16xf16>, vector<16xf16> + } else { + %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>) + %generic_gate, %generic_up = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_generic_block>(%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, %quant_block, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_weight, %up_weight, %expert_byte_base, %gate_weight_row_bytes, %up_weight_row_bytes, %gate_weight_format, %up_weight_format, %gate_accumulator, %up_accumulator) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, offset, index, index, vector<16xf16>, vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + scf.yield %generic_gate, %generic_up : vector<16xf16>, vector<16xf16> + } + %q4_or_generic_gate = scf.select %is_q4k, %q4_gate, %generic_gate : vector<16xf16> + %q4_or_generic_up = scf.select %is_q4k, %q4_up, %generic_up : vector<16xf16> + %q5_or_prior_gate = scf.select %is_q5k, %q5_gate, %q4_or_generic_gate : vector<16xf16> + %q5_or_prior_up = scf.select %is_q5k, %q5_up, %q4_or_generic_up : vector<16xf16> + %next_gate = scf.select %is_q6k, %q6_gate, %q5_or_prior_gate : vector<16xf16> + %next_up = scf.select %is_q6k, %q6_up, %q5_or_prior_up : vector<16xf16> + template.return %next_gate, %next_up : vector<16xf16>, vector<16xf16> +} +// Generic routed gate/up SwiGLU accumulation for FP16 WMMA. +// +// This path covers formats whose values can be staged directly through the +// shared dequant helper. Q4_K and Q6_K keep specialized retained-packet paths. + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.project_generic_vector16(%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>, %stage_activation: i1, %weight_format: index, %quant_block: index, %quant_group: index, %k_origin: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %accumulator: vector<16xf16>) -> (vector<16xf16>) + +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>) + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.project_generic_vector16> device @ggml_mul_mat_id_swiglu_f16_f16_project_generic_vector16(%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>, %stage_activation: i1, %weight_format: index, %quant_block: index, %quant_group: index, %k_origin: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %accumulator: vector<16xf16>) -> (vector<16xf16>) { + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %input_noalias, %route_stage_noalias, %weight_stage_noalias, %activation_stage_noalias, %weight_noalias = buffer.assume.noalias %input, %route_stage, %weight_stage, %activation_stage, %weight : buffer, buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> + %route_stage_view = buffer.view %route_stage_noalias[%c0_offset] : buffer -> view<32xi32> + %weight_stage_view = buffer.view %weight_stage_noalias[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %bounded_group = index.assume %quant_group [range(%quant_group, 0, 7)] : 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_values = scf.if %valid_channel -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %weight_k = index.add %k_origin, %load_k : index + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %input_size, %quant_block, %bounded_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> + } + scf.if %stage_activation { + %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %route_count : index + %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next_accumulator = scf.for %k_half = [%c0 to %c32 step %c16](%half_accumulator = %accumulator : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %next = vector.mma %lhs, %rhs, %half_accumulator : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return %next_accumulator : vector<16xf16> +} + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_generic_block> device @ggml_mul_mat_id_swiglu_f16_f16_accumulate_generic_block(%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>, %quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %gate_weight_row_bytes: offset, %up_weight_row_bytes: offset, %gate_weight_format: index, %up_weight_format: index, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %true = scalar.constant true : i1 + %false = scalar.constant false : i1 + %next_block_gate, %next_block_up = scf.for %quant_group = [%c0 to %c8 step %c1](%group_gate = %gate_accumulator : vector<16xf16>, %group_up = %up_accumulator : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %bounded_group = index.assume %quant_group [range(%quant_group, 0, 7)] : index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %bounded_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + %gate_next = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.project_generic_vector16>(%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, %true, %gate_weight_format, %quant_block, %bounded_group, %k_origin, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_weight, %expert_byte_base, %gate_weight_row_bytes, %group_gate) : (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>, i1, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<16xf16>) -> (vector<16xf16>) + %up_next = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.project_generic_vector16>(%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, %false, %up_weight_format, %quant_block, %bounded_group, %k_origin, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %up_weight, %expert_byte_base, %up_weight_row_bytes, %group_up) : (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>, i1, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<16xf16>) -> (vector<16xf16>) + scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> + } + template.return %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> +} +// Q4_K routed gate/up SwiGLU accumulation. +// +// The shared body owns routing and publish. This file only binds Q4_K block +// accumulation to that common schedule. + +template.decl @ggml.mul_mat_id_q4k_f16_wmma.project_vector16(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: index, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: vector<4xi32>, %arg18: vector<4xi32>, %arg19: vector<1xi32>, %arg20: vector<1xi32>, %arg21: vector<16xf16>) -> (vector<16xf16>) + +func.decl @ggml_q4k_split_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) + +func.decl @ggml_q4k_split_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_q4k_block> device @ggml_mul_mat_id_swiglu_f16_f16_accumulate_q4k_block(%q4_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_i32x1 = vector.constant 0 : vector<1xi32> + %c0_i32x4 = vector.constant 0 : vector<4xi32> + %true = scalar.constant true : i1 + %false = scalar.constant false : i1 + %gate_header0, %gate_header1, %up_header0, %up_header1 = scf.for %header_row_offset = [%c0 to %c64 step %c32](%prior_gate_header0 = %c0_i32x4 : vector<4xi32>, %prior_gate_header1 = %c0_i32x4 : vector<4xi32>, %prior_up_header0 = %c0_i32x4 : vector<4xi32>, %prior_up_header1 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %header_local_row0 = index.add %load_row, %header_row_offset : index + %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index + %header_channel = index.add %channel_tile_base, %header_local_row : index + %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index + %loaded_gate_header, %loaded_up_header = scf.if %valid_header_channel -> (vector<4xi32>, vector<4xi32>) { + %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset + %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset + %gate_header = func.call @ggml_q4k_split_wmma_load_header(%gate_weight, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) + %up_header = func.call @ggml_q4k_split_wmma_load_header(%up_weight, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) + scf.yield %gate_header, %up_header : vector<4xi32>, vector<4xi32> + } else { + scf.yield %c0_i32x4, %c0_i32x4 : vector<4xi32>, vector<4xi32> + } + %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index + %next_gate_header0 = scf.select %updates_header0, %loaded_gate_header, %prior_gate_header0 : vector<4xi32> + %next_gate_header1 = scf.select %updates_header0, %prior_gate_header1, %loaded_gate_header : vector<4xi32> + %next_up_header0 = scf.select %updates_header0, %loaded_up_header, %prior_up_header0 : vector<4xi32> + %next_up_header1 = scf.select %updates_header0, %prior_up_header1, %loaded_up_header : vector<4xi32> + scf.yield %next_gate_header0, %next_gate_header1, %next_up_header0, %next_up_header1 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %next_block_gate, %next_block_up = scf.for %q4_group_pair = [%c0 to %c4 step %c1](%pair_gate_acc = %gate_accumulator : vector<16xf16>, %pair_up_acc = %up_accumulator : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %gate_q_word0, %gate_q_word1, %up_q_word0, %up_q_word1 = scf.for %code_row_offset = [%c0 to %c64 step %c32](%prior_gate_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word1 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %code_local_row0 = index.add %load_row, %code_row_offset : index + %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index + %code_channel = index.add %channel_tile_base, %code_local_row : index + %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index + %loaded_gate_q_word, %loaded_up_q_word = scf.if %valid_code_channel -> (vector<1xi32>, vector<1xi32>) { + %bounded_q4_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset + %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset + %gate_q_word = func.call @ggml_q4k_split_wmma_load_code(%gate_weight, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + %up_q_word = func.call @ggml_q4k_split_wmma_load_code(%up_weight, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + scf.yield %gate_q_word, %up_q_word : vector<1xi32>, vector<1xi32> + } else { + scf.yield %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32> + } + %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index + %next_gate_q_word0 = scf.select %updates_q_word0, %loaded_gate_q_word, %prior_gate_q_word0 : vector<1xi32> + %next_gate_q_word1 = scf.select %updates_q_word0, %prior_gate_q_word1, %loaded_gate_q_word : vector<1xi32> + %next_up_q_word0 = scf.select %updates_q_word0, %loaded_up_q_word, %prior_up_q_word0 : vector<1xi32> + %next_up_q_word1 = scf.select %updates_q_word0, %prior_up_q_word1, %loaded_up_q_word : vector<1xi32> + scf.yield %next_gate_q_word0, %next_gate_q_word1, %next_up_q_word0, %next_up_q_word1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %next_pair_gate, %next_pair_up = scf.for %group_within_pair = [%c0 to %c2 step %c1](%group_gate = %pair_gate_acc : vector<16xf16>, %group_up = %pair_up_acc : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %q4_group_base = index.mul %q4_group_pair, %c2 : index + %q4_group0 = index.add %q4_group_base, %group_within_pair : index + %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index + %block_k_base = index.mul %q4_block, %c256 : index + %group_k_add = index.mul %q4_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + %gate_next = template.apply<@ggml.mul_mat_id_q4k_f16_wmma.project_vector16>(%true, %q4_group, %k_origin, %load_row, %load_k, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_header0, %gate_header1, %gate_q_word0, %gate_q_word1, %group_gate) : (i1, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, vector<4xi32>, vector<4xi32>, vector<1xi32>, vector<1xi32>, vector<16xf16>) -> (vector<16xf16>) + %up_next = template.apply<@ggml.mul_mat_id_q4k_f16_wmma.project_vector16>(%false, %q4_group, %k_origin, %load_row, %load_k, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %up_header0, %up_header1, %up_q_word0, %up_q_word1, %group_up) : (i1, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, vector<4xi32>, vector<4xi32>, vector<1xi32>, vector<1xi32>, vector<16xf16>) -> (vector<16xf16>) + scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> + } + scf.yield %next_pair_gate, %next_pair_up : vector<16xf16>, vector<16xf16> + } + template.return %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> +} +// Q5_K routed gate/up SwiGLU accumulation. +// +// Q5_K reuses the Q4_K group-pair low-bit schedule and additionally retains the +// per-packet high-bit bytes across all groups in the block. + +template.decl @ggml.mul_mat_id_q5k_f16_wmma.project_vector16(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: index, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: vector<4xi32>, %arg18: vector<4xi32>, %arg19: vector<1xi32>, %arg20: vector<1xi32>, %arg21: vector<4xi8>, %arg22: vector<4xi8>, %arg23: vector<16xf16>) -> (vector<16xf16>) + +func.decl @ggml_q5k_split_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q5_block: index) -> (vector<4xi32>) + +func.decl @ggml_q5k_split_wmma_load_high_bits(%weight: buffer, %row_byte_base: offset, %q5_block: index, %packet: index) -> (vector<4xi8>) + +func.decl @ggml_q5k_split_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group_pair: index, %packet: index) -> (vector<1xi32>) + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_q5k_block> device @ggml_mul_mat_id_swiglu_f16_f16_accumulate_q5k_block(%q5_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_i32x1 = vector.constant 0 : vector<1xi32> + %c0_i32x4 = vector.constant 0 : vector<4xi32> + %c0_i8x4 = vector.constant 0 : vector<4xi8> + %true = scalar.constant true : i1 + %false = scalar.constant false : i1 + %gate_header0, %gate_header1, %up_header0, %up_header1, %gate_qh_bytes0, %gate_qh_bytes1, %up_qh_bytes0, %up_qh_bytes1 = scf.for %header_row_offset = [%c0 to %c64 step %c32](%prior_gate_header0 = %c0_i32x4 : vector<4xi32>, %prior_gate_header1 = %c0_i32x4 : vector<4xi32>, %prior_up_header0 = %c0_i32x4 : vector<4xi32>, %prior_up_header1 = %c0_i32x4 : vector<4xi32>, %prior_gate_qh_bytes0 = %c0_i8x4 : vector<4xi8>, %prior_gate_qh_bytes1 = %c0_i8x4 : vector<4xi8>, %prior_up_qh_bytes0 = %c0_i8x4 : vector<4xi8>, %prior_up_qh_bytes1 = %c0_i8x4 : vector<4xi8>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) unroll { + %header_local_row0 = index.add %load_row, %header_row_offset : index + %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index + %header_channel = index.add %channel_tile_base, %header_local_row : index + %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index + %loaded_gate_header, %loaded_up_header, %loaded_gate_qh_bytes, %loaded_up_qh_bytes = scf.if %valid_header_channel -> (vector<4xi32>, vector<4xi32>, vector<4xi8>, vector<4xi8>) { + %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset + %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset + %gate_header = func.call @ggml_q5k_split_wmma_load_header(%gate_weight, %header_row_byte_base, %q5_block) : (buffer, offset, index) -> (vector<4xi32>) + %up_header = func.call @ggml_q5k_split_wmma_load_header(%up_weight, %header_row_byte_base, %q5_block) : (buffer, offset, index) -> (vector<4xi32>) + %gate_qh_bytes = func.call @ggml_q5k_split_wmma_load_high_bits(%gate_weight, %header_row_byte_base, %q5_block, %load_packet) : (buffer, offset, index, index) -> (vector<4xi8>) + %up_qh_bytes = func.call @ggml_q5k_split_wmma_load_high_bits(%up_weight, %header_row_byte_base, %q5_block, %load_packet) : (buffer, offset, index, index) -> (vector<4xi8>) + scf.yield %gate_header, %up_header, %gate_qh_bytes, %up_qh_bytes : vector<4xi32>, vector<4xi32>, vector<4xi8>, vector<4xi8> + } else { + scf.yield %c0_i32x4, %c0_i32x4, %c0_i8x4, %c0_i8x4 : vector<4xi32>, vector<4xi32>, vector<4xi8>, vector<4xi8> + } + %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index + %next_gate_header0 = scf.select %updates_header0, %loaded_gate_header, %prior_gate_header0 : vector<4xi32> + %next_gate_header1 = scf.select %updates_header0, %prior_gate_header1, %loaded_gate_header : vector<4xi32> + %next_up_header0 = scf.select %updates_header0, %loaded_up_header, %prior_up_header0 : vector<4xi32> + %next_up_header1 = scf.select %updates_header0, %prior_up_header1, %loaded_up_header : vector<4xi32> + %next_gate_qh_bytes0 = scf.select %updates_header0, %loaded_gate_qh_bytes, %prior_gate_qh_bytes0 : vector<4xi8> + %next_gate_qh_bytes1 = scf.select %updates_header0, %prior_gate_qh_bytes1, %loaded_gate_qh_bytes : vector<4xi8> + %next_up_qh_bytes0 = scf.select %updates_header0, %loaded_up_qh_bytes, %prior_up_qh_bytes0 : vector<4xi8> + %next_up_qh_bytes1 = scf.select %updates_header0, %prior_up_qh_bytes1, %loaded_up_qh_bytes : vector<4xi8> + scf.yield %next_gate_header0, %next_gate_header1, %next_up_header0, %next_up_header1, %next_gate_qh_bytes0, %next_gate_qh_bytes1, %next_up_qh_bytes0, %next_up_qh_bytes1 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8> + } + %next_block_gate, %next_block_up = scf.for %q5_group_pair = [%c0 to %c4 step %c1](%pair_gate_acc = %gate_accumulator : vector<16xf16>, %pair_up_acc = %up_accumulator : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %gate_q_word0, %gate_q_word1, %up_q_word0, %up_q_word1 = scf.for %code_row_offset = [%c0 to %c64 step %c32](%prior_gate_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word1 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %code_local_row0 = index.add %load_row, %code_row_offset : index + %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index + %code_channel = index.add %channel_tile_base, %code_local_row : index + %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index + %loaded_gate_q_word, %loaded_up_q_word = scf.if %valid_code_channel -> (vector<1xi32>, vector<1xi32>) { + %bounded_q5_group_pair = index.assume %q5_group_pair [range(%q5_group_pair, 0, 3)] : index + %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset + %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset + %gate_q_word = func.call @ggml_q5k_split_wmma_load_code(%gate_weight, %code_row_byte_base, %q5_block, %bounded_q5_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + %up_q_word = func.call @ggml_q5k_split_wmma_load_code(%up_weight, %code_row_byte_base, %q5_block, %bounded_q5_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + scf.yield %gate_q_word, %up_q_word : vector<1xi32>, vector<1xi32> + } else { + scf.yield %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32> + } + %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index + %next_gate_q_word0 = scf.select %updates_q_word0, %loaded_gate_q_word, %prior_gate_q_word0 : vector<1xi32> + %next_gate_q_word1 = scf.select %updates_q_word0, %prior_gate_q_word1, %loaded_gate_q_word : vector<1xi32> + %next_up_q_word0 = scf.select %updates_q_word0, %loaded_up_q_word, %prior_up_q_word0 : vector<1xi32> + %next_up_q_word1 = scf.select %updates_q_word0, %prior_up_q_word1, %loaded_up_q_word : vector<1xi32> + scf.yield %next_gate_q_word0, %next_gate_q_word1, %next_up_q_word0, %next_up_q_word1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %next_pair_gate, %next_pair_up = scf.for %group_within_pair = [%c0 to %c2 step %c1](%group_gate = %pair_gate_acc : vector<16xf16>, %group_up = %pair_up_acc : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %q5_group_base = index.mul %q5_group_pair, %c2 : index + %q5_group0 = index.add %q5_group_base, %group_within_pair : index + %q5_group = index.assume %q5_group0 [range(%q5_group0, 0, 7)] : index + %block_k_base = index.mul %q5_block, %c256 : index + %group_k_add = index.mul %q5_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + %gate_next = template.apply<@ggml.mul_mat_id_q5k_f16_wmma.project_vector16>(%true, %q5_group, %k_origin, %load_row, %load_k, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_header0, %gate_header1, %gate_q_word0, %gate_q_word1, %gate_qh_bytes0, %gate_qh_bytes1, %group_gate) : (i1, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, vector<4xi32>, vector<4xi32>, vector<1xi32>, vector<1xi32>, vector<4xi8>, vector<4xi8>, vector<16xf16>) -> (vector<16xf16>) + %up_next = template.apply<@ggml.mul_mat_id_q5k_f16_wmma.project_vector16>(%false, %q5_group, %k_origin, %load_row, %load_k, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %up_header0, %up_header1, %up_q_word0, %up_q_word1, %up_qh_bytes0, %up_qh_bytes1, %group_up) : (i1, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, vector<4xi32>, vector<4xi32>, vector<1xi32>, vector<1xi32>, vector<4xi8>, vector<4xi8>, vector<16xf16>) -> (vector<16xf16>) + scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> + } + scf.yield %next_pair_gate, %next_pair_up : vector<16xf16>, vector<16xf16> + } + template.return %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> +} +// Q6_K routed gate/up SwiGLU accumulation. +// +// The shared body owns routing and publish. This file only binds Q6_K block +// accumulation to that common schedule. + +template.decl @ggml.mul_mat_id_q6k_f16_wmma.project_vector16(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: index, %arg13: index, %arg14: index, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: offset, %arg21: offset, %arg22: vector<1xi32>, %arg23: vector<1xi32>, %arg24: vector<1xi32>, %arg25: vector<1xi32>, %arg26: vector<1xi32>, %arg27: vector<1xi32>, %arg28: vector<16xf16>) -> (vector<16xf16>) + +func.decl @ggml_q6k_split_wmma_load_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_half: index, %packet: index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_q6k_block> device @ggml_mul_mat_id_swiglu_f16_f16_accumulate_q6k_block(%q6_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_i32x1 = vector.constant 0 : vector<1xi32> + %true = scalar.constant true : i1 + %false = scalar.constant false : i1 + %next_block_gate, %next_block_up = scf.for %q6_half = [%c0 to %c2 step %c1](%half_gate_acc = %gate_accumulator : vector<16xf16>, %half_up_acc = %up_accumulator : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %gate_ql0_word0, %gate_ql1_word0, %gate_qh_word0, %gate_ql0_word1, %gate_ql1_word1, %gate_qh_word1, %up_ql0_word0, %up_ql1_word0, %up_qh_word0, %up_ql0_word1, %up_ql1_word1, %up_qh_word1 = scf.for %code_row_offset = [%c0 to %c64 step %c32](%prior_gate_ql0_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_ql1_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_qh_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_ql0_word1 = %c0_i32x1 : vector<1xi32>, %prior_gate_ql1_word1 = %c0_i32x1 : vector<1xi32>, %prior_gate_qh_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_ql0_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_ql1_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_qh_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_ql0_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_ql1_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_qh_word1 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { + %code_local_row0 = index.add %load_row, %code_row_offset : index + %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index + %code_channel = index.add %channel_tile_base, %code_local_row : index + %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index + %loaded_gate_ql0, %loaded_gate_ql1, %loaded_gate_qh, %loaded_up_ql0, %loaded_up_ql1, %loaded_up_qh = scf.if %valid_code_channel -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) { + %bounded_q6_half = index.assume %q6_half [range(%q6_half, 0, 1)] : index + %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset + %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset + %gate_ql0, %gate_ql1, %gate_qh = func.call @ggml_q6k_split_wmma_load_half_codes(%gate_weight, %code_row_byte_base, %q6_block, %bounded_q6_half, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) + %up_ql0, %up_ql1, %up_qh = func.call @ggml_q6k_split_wmma_load_half_codes(%up_weight, %code_row_byte_base, %q6_block, %bounded_q6_half, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) + scf.yield %gate_ql0, %gate_ql1, %gate_qh, %up_ql0, %up_ql1, %up_qh : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } else { + scf.yield %c0_i32x1, %c0_i32x1, %c0_i32x1, %c0_i32x1, %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %updates_word0 = index.cmp eq, %code_row_offset, %c0 : index + %next_gate_ql0_word0 = scf.select %updates_word0, %loaded_gate_ql0, %prior_gate_ql0_word0 : vector<1xi32> + %next_gate_ql1_word0 = scf.select %updates_word0, %loaded_gate_ql1, %prior_gate_ql1_word0 : vector<1xi32> + %next_gate_qh_word0 = scf.select %updates_word0, %loaded_gate_qh, %prior_gate_qh_word0 : vector<1xi32> + %next_gate_ql0_word1 = scf.select %updates_word0, %prior_gate_ql0_word1, %loaded_gate_ql0 : vector<1xi32> + %next_gate_ql1_word1 = scf.select %updates_word0, %prior_gate_ql1_word1, %loaded_gate_ql1 : vector<1xi32> + %next_gate_qh_word1 = scf.select %updates_word0, %prior_gate_qh_word1, %loaded_gate_qh : vector<1xi32> + %next_up_ql0_word0 = scf.select %updates_word0, %loaded_up_ql0, %prior_up_ql0_word0 : vector<1xi32> + %next_up_ql1_word0 = scf.select %updates_word0, %loaded_up_ql1, %prior_up_ql1_word0 : vector<1xi32> + %next_up_qh_word0 = scf.select %updates_word0, %loaded_up_qh, %prior_up_qh_word0 : vector<1xi32> + %next_up_ql0_word1 = scf.select %updates_word0, %prior_up_ql0_word1, %loaded_up_ql0 : vector<1xi32> + %next_up_ql1_word1 = scf.select %updates_word0, %prior_up_ql1_word1, %loaded_up_ql1 : vector<1xi32> + %next_up_qh_word1 = scf.select %updates_word0, %prior_up_qh_word1, %loaded_up_qh : vector<1xi32> + scf.yield %next_gate_ql0_word0, %next_gate_ql1_word0, %next_gate_qh_word0, %next_gate_ql0_word1, %next_gate_ql1_word1, %next_gate_qh_word1, %next_up_ql0_word0, %next_up_ql1_word0, %next_up_qh_word0, %next_up_ql0_word1, %next_up_ql1_word1, %next_up_qh_word1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> + } + %next_half_gate, %next_half_up = scf.for %group_within_half = [%c0 to %c4 step %c1](%group_gate = %half_gate_acc : vector<16xf16>, %group_up = %half_up_acc : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %q6_group_base = index.mul %q6_half, %c4 : index + %q6_group0 = index.add %q6_group_base, %group_within_half : index + %q6_group = index.assume %q6_group0 [range(%q6_group0, 0, 7)] : index + %block_k_base = index.mul %q6_block, %c256 : index + %group_k_add = index.mul %q6_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + %gate_next = template.apply<@ggml.mul_mat_id_q6k_f16_wmma.project_vector16>(%true, %q6_group, %q6_block, %k_origin, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %gate_weight, %expert_byte_base, %weight_row_bytes, %gate_ql0_word0, %gate_ql1_word0, %gate_qh_word0, %gate_ql0_word1, %gate_ql1_word1, %gate_qh_word1, %group_gate) : (i1, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<16xf16>) -> (vector<16xf16>) + %up_next = template.apply<@ggml.mul_mat_id_q6k_f16_wmma.project_vector16>(%false, %q6_group, %q6_block, %k_origin, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input, %route_stage, %weight_stage, %activation_stage, %up_weight, %expert_byte_base, %weight_row_bytes, %up_ql0_word0, %up_ql1_word0, %up_qh_word0, %up_ql0_word1, %up_ql1_word1, %up_qh_word1, %group_up) : (i1, index, index, index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, offset, offset, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<16xf16>) -> (vector<16xf16>) + scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> + } + scf.yield %next_half_gate, %next_half_up : vector<16xf16>, vector<16xf16> + } + template.return %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_binary_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_binary_publish.loom new file mode 100644 index 000000000000..c2a9bb0583d6 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_binary_publish.loom @@ -0,0 +1,14 @@ +// SwiGLU activation and F16 publication for routed pair results. +template.decl @ggml.mul_mat_id_swiglu_f16_f16.binary_publish_vector8(%gate: vector<8xf32>, %up: vector<8xf32>, %assignment: index, %channel: index, %output_row_count: index, %output_size: index, %output: buffer) + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.binary_publish_vector8> device @ggml_mul_mat_id_swiglu_f16_f16_binary_publish_vector8(%gate: vector<8xf32>, %up: vector<8xf32>, %assignment: index, %channel: index, %output_row_count: index, %output_size: index, %output: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %activated = vector.siluf %gate : vector<8xf32> + %wide_values = vector.mulf %activated, %up : vector<8xf32> + %values = vector.fptrunc %wide_values : vector<8xf32> to vector<8xf16> + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%output_row_count]x[%output_size]xf16> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<8xi1> + vector.store.mask %values, %output_view[%assignment, %channel], %mask : vector<8xf16>, view<[%output_row_count]x[%output_size]xf16>, vector<8xi1> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_schedule.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_schedule.loom new file mode 100644 index 000000000000..3662259a5f3c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_pair_schedule.loom @@ -0,0 +1,183 @@ +// Shared routed gate/up schedule for F16 pair WMMA fast paths. +// +// The body owns routing, staging lifetime, and accumulator flow. Separate +// motifs select the dtype-specific inner loop and publish the binary result. + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.accumulate_block(%quant_block: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %gate_weight: buffer, %up_weight: buffer, %expert_byte_base: offset, %gate_weight_row_bytes: offset, %up_weight_row_bytes: offset, %gate_weight_format: index, %up_weight_format: index, %gate_accumulator: vector<16xf16>, %up_accumulator: vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %route_count: index, %expert_count: index, %output_size: index, %descriptor_expert_mask: i32, %descriptor_partition_shift: i32, %descriptor_row_count_shift: i32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.binary_publish_vector8(%gate: vector<8xf32>, %up: vector<8xf32>, %assignment: index, %channel: index, %output_row_count: index, %output_size: index, %output: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.def inline @ggml_mul_mat_id_swiglu_f16_f16_unpack_expert_partition_descriptor(%descriptor: i32, %expert_mask: i32, %partition_shift: i32, %row_count_shift: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %expert_i32 = scalar.andi %descriptor, %expert_mask : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %partition_shift : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %row_count_shift : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +template.def<@ggml.mul_mat_id_swiglu_f16_f16.body> device @ggml_mul_mat_id_swiglu_f16_f16_wmma_body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %route_count: index, %expert_count: index, %output_size: index, %descriptor_expert_mask: i32, %descriptor_partition_shift: i32, %descriptor_row_count_shift: i32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 4096 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %zero_accumulator = vector.constant 0.0 : vector<16xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %assignment_count = index.mul %token_count, %bounded_route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %gate_weight_block_bytes = func.call @ggml_dequant_weight_tile_bytes(%gate_weight_format) : (index) -> (offset) + %up_weight_block_bytes = func.call @ggml_dequant_weight_tile_bytes(%up_weight_format) : (index) -> (offset) + %quant_block_count = index.div %input_size, %c256 : index + %gate_weight_row_bytes = index.scale %quant_block_count, %gate_weight_block_bytes : index, offset -> offset + %up_weight_row_bytes = index.scale %quant_block_count, %up_weight_block_bytes : index, offset -> offset + %gate_weight_expert_bytes = index.scale %bounded_output_size, %gate_weight_row_bytes : index, offset -> offset + %output_row_count = index.mul %token_count, %bounded_route_count : index + %input_noalias, %expert_table_noalias, %partition_table_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [lt(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @ggml_mul_mat_id_swiglu_f16_f16_unpack_expert_partition_descriptor(%descriptor_i32, %descriptor_expert_mask, %descriptor_partition_shift, %descriptor_row_count_shift) : (i32, i32, i32, i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 31)] : index + %expert_byte_base = index.scale %bounded_expert, %gate_weight_expert_bytes : index, offset -> offset + %channel_subgroup = index.rem %subgroup, %c4 : index + %route_subgroup = index.div %subgroup, %c4 : index + %subgroup_channel_add = index.mul %channel_subgroup, %c16 : index + %subgroup_route_add = index.mul %route_subgroup, %c16 : index + %init_gate = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %init_up = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %gate_result, %up_result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_gate = %init_gate : vector<16xf16>, %block_up = %init_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { + %next_gate, %next_up = template.apply<@ggml.mul_mat_id_swiglu_f16_f16.accumulate_block>(%quant_block, %load_row, %load_k, %load_packet, %channel_tile_base, %bounded_output_size, %assignment_count, %token_count, %bounded_route_count, %input_size, %subgroup_channel_add, %subgroup_route_add, %input_noalias, %route_stage, %weight_stage, %activation_stage, %gate_weight_noalias, %up_weight_noalias, %expert_byte_base, %gate_weight_row_bytes, %up_weight_row_bytes, %gate_weight_format, %up_weight_format, %block_gate, %block_up) : (index, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, offset, index, index, vector<16xf16>, vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) + scf.yield %next_gate, %next_up : vector<16xf16>, vector<16xf16> + } + %publish_route0 = index.div %lane, %c2 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c2 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index + %publish_channel_add = index.mul %publish_packet, %c8 : index + %local_route = index.add %subgroup_route_add, %publish_route : index + %assignment_i32 = view.load %route_stage_view[%local_route] : view<32xi32> -> i32 + %assignment_nonnegative = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %safe_assignment_i32 = scf.if %assignment_nonnegative -> (i32) { + scf.yield %assignment_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0 = index.cast %safe_assignment_i32 : i32 to index + %safe_assignment = index.assume %safe_assignment0 [range(%safe_assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %safe_assignment, %assignment_count [lt(%safe_assignment, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel = index.add %subgroup_channel_base, %publish_channel_add : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %writes = scalar.andi %assignment_nonnegative, %valid_channel : i1 + vector.fragment.store %gate_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + %gate_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %gate_wide = vector.extf %gate_values : vector<8xf16> to vector<8xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %up_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes { + %up_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %up_wide = vector.extf %up_values : vector<8xf16> to vector<8xf32> + template.apply<@ggml.mul_mat_id_swiglu_f16_f16.binary_publish_vector8>(%gate_wide, %up_wide, %bounded_assignment, %channel, %output_row_count, %bounded_output_size, %output_noalias) : (vector<8xf32>, vector<8xf32>, index, index, index, index, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q4k.loom new file mode 100644 index 000000000000..a429c387797d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q4k.loom @@ -0,0 +1,200 @@ +// Shared routed F16 pair Q4_K single-projection WMMA step. +// +// Callers own the outer routing and fusion schedule. This motif stages one +// projected Q4_K weight tile, optionally stages the shared routed activation +// tile, performs the 16x16 WMMA update, and returns the updated accumulator. + +func.def inline @ggml_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +// Acquires the packed code word shared by one adjacent Q4_K group pair. +func.def inline @ggml_q4k_split_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) { + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %bounded_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %q_page = index.mul %bounded_group_pair, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + func.return %q_word : vector<1xi32> +} + +// Decodes the four adjacent Q4_K values owned by one load packet from an +// already-loaded block header and packed code word. Matrix schedules choose +// the lifetime of both immutable packets. +func.def inline @ggml_q4k_split_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %scale, %minimum = func.call @ggml_q4k_scale_from_header(%scale0, %scale1, %scale2, %bounded_group) : (i32, i32, i32, index) -> (i32, i32) + %scale_f32 = scalar.uitofp %scale : i32 to f32 + %minimum_f32 = scalar.uitofp %minimum : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + %minimum_scale = scalar.mulf %dmin, %minimum_f32 : f32 + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + // Form adjacent FP16 lanes from fused FP32 affine expressions. AMDGPU maps + // this natural shape to packed mixlo/mixhi instructions where available. + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %half0 = scalar.fptrunc %value0 : f32 to f16 + %half1 = scalar.fptrunc %value1 : f32 to f16 + %half2 = scalar.fptrunc %value2 : f32 to f16 + %half3 = scalar.fptrunc %value3 : f32 to f16 + %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> + func.return %result : vector<4xf16> +} + +// Decodes one group when its caller has retained only the Q4_K block header. +func.def inline @ggml_q4k_split_wmma_vector4_from_header(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index, %header_words: vector<4xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %q4_group_pair = index.div %bounded_group, %c2 : index + %q_word = func.call @ggml_q4k_split_wmma_load_code(%weight, %row_byte_base, %q4_block, %q4_group_pair, %packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) + %values = func.call @ggml_q4k_split_wmma_vector4_from_header_code(%bounded_group, %header_words, %q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) + func.return %values : vector<4xf16> +} + +// Acquires one naturally aligned Q4_K block header as a single 16-byte packet. +func.def inline @ggml_q4k_split_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) { + %c0 = index.constant 0 : index + %block_bytes = index.constant 144 : offset + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> + %header_words = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> + func.return %header_words : vector<4xi32> +} + +// Acquires one block header before decoding the selected four-value group +// packet. Matrix schedules that span several groups call the two operations +// separately so the header lifetime matches their complete block loop. +func.def inline @ggml_q4k_split_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { + %header_words = func.call @ggml_q4k_split_wmma_load_header(%weight, %row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) + %values = func.call @ggml_q4k_split_wmma_vector4_from_header(%weight, %row_byte_base, %q4_block, %q4_group, %packet, %header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<4xf16>) + func.return %values : vector<4xf16> +} + +template.decl @ggml.mul_mat_id_q4k_f16_wmma.project_vector16(%stage_activation: i1, %q4_group: index, %k_origin: index, %load_row: index, %load_k: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %header0: vector<4xi32>, %header1: vector<4xi32>, %q_word0: vector<1xi32>, %q_word1: vector<1xi32>, %accumulator: vector<16xf16>) -> (vector<16xf16>) + +template.def<@ggml.mul_mat_id_q4k_f16_wmma.project_vector16> device @ggml_mul_mat_id_q4k_f16_wmma_project_vector16(%stage_activation: i1, %q4_group: index, %k_origin: index, %load_row: index, %load_k: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %header0: vector<4xi32>, %header1: vector<4xi32>, %q_word0: vector<1xi32>, %q_word1: vector<1xi32>, %accumulator: vector<16xf16>) -> (vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %input_noalias, %route_stage_noalias, %weight_stage_noalias, %activation_stage_noalias = buffer.assume.noalias %input, %route_stage, %weight_stage, %activation_stage : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> + %route_stage_view = buffer.view %route_stage_noalias[%c0_offset] : buffer -> view<32xi32> + %weight_stage_view = buffer.view %weight_stage_noalias[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : 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 + %selects_first_row = index.cmp eq, %row_offset, %c0 : index + %selected_header = scf.select %selects_first_row, %header0, %header1 : vector<4xi32> + %selected_q_word = scf.select %selects_first_row, %q_word0, %q_word1 : vector<1xi32> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %decoded = func.call @ggml_q4k_split_wmma_vector4_from_header_code(%bounded_group, %selected_header, %selected_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (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> + } + scf.if %stage_activation { + %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %route_count : index + %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next_accumulator = scf.for %k_half = [%c0 to %c32 step %c16](%half_accumulator = %accumulator : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %next = vector.mma %lhs, %rhs, %half_accumulator : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return %next_accumulator : vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q5k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q5k.loom new file mode 100644 index 000000000000..14042cfb38c1 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q5k.loom @@ -0,0 +1,221 @@ +// Shared routed F16 pair Q5_K single-projection WMMA step. +// +// Callers own the outer routing and fusion schedule. This motif stages one +// projected Q5_K weight tile, optionally stages the shared routed activation +// tile, performs the 16x16 WMMA update, and returns the updated accumulator. + +func.def inline @ggml_q5k_split_wmma_high_bit_f32(%qh: i8, %qh_bit: i32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c16_i32 = scalar.constant 16 : i32 + %qh_i32 = scalar.extui %qh : i8 to i32 + %masked = scalar.andi %qh_i32, %qh_bit : i32 + %is_set = scalar.cmpi ne, %masked, %c0_i32 : i32 + %high = scf.select %is_set, %c16_i32, %c0_i32 : i32 + %high_f32 = scalar.uitofp %high : i32 to f32 + func.return %high_f32 : f32 +} + +func.def inline @ggml_q5k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q5_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q5_group [range(%q5_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +// Acquires the packed low-bit code word shared by one adjacent Q5_K group pair. +func.def inline @ggml_q5k_split_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group_pair: index, %packet: index) -> (vector<1xi32>) { + %c8 = index.constant 8 : index + %block_bytes = index.constant 176 : offset + %code_offset = index.constant 48 : offset + %bounded_group_pair = index.assume %q5_group_pair [range(%q5_group_pair, 0, 3)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q5_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %q_page = index.mul %bounded_group_pair, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + func.return %q_word : vector<1xi32> +} + +// Acquires the four Q5_K high-bit bytes owned by one packet. The same bytes are +// reused by all eight groups in the block. +func.def inline @ggml_q5k_split_wmma_load_high_bits(%weight: buffer, %row_byte_base: offset, %q5_block: index, %packet: index) -> (vector<4xi8>) { + %c4 = index.constant 4 : index + %block_bytes = index.constant 176 : offset + %qh_offset = index.constant 16 : offset + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q5_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 + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<32xi8> + %packet_base = index.mul %bounded_packet, %c4 : index + %qh_bytes = vector.load %qh_view[%packet_base] : view<32xi8> -> vector<4xi8> + func.return %qh_bytes : vector<4xi8> +} + +// Decodes the four adjacent Q5_K values owned by one load packet from retained +// block metadata and packed code. +func.def inline @ggml_q5k_split_wmma_vector4_from_header_code(%q5_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>, %qh_bytes: vector<4xi8>) -> (vector<4xf16>) { + %c1_i32 = scalar.constant 1 : i32 + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q5_group [range(%q5_group, 0, 7)] : index + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %scale, %minimum = func.call @ggml_q5k_scale_from_header(%scale0, %scale1, %scale2, %bounded_group) : (i32, i32, i32, index) -> (i32, i32) + %scale_f32 = scalar.uitofp %scale : i32 to f32 + %minimum_f32 = scalar.uitofp %minimum : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + %minimum_scale = scalar.mulf %dmin, %minimum_f32 : f32 + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %qh_shift_i32 = index.cast %bounded_group : index to i32 + %qh_bit = scalar.shli %c1_i32, %qh_shift_i32 : i32 + %qh0 = vector.extract %qh_bytes[0] : vector<4xi8> -> i8 + %qh1 = vector.extract %qh_bytes[1] : vector<4xi8> -> i8 + %qh2 = vector.extract %qh_bytes[2] : vector<4xi8> -> i8 + %qh3 = vector.extract %qh_bytes[3] : vector<4xi8> -> i8 + %high0 = func.call @ggml_q5k_split_wmma_high_bit_f32(%qh0, %qh_bit) : (i8, i32) -> (f32) + %high1 = func.call @ggml_q5k_split_wmma_high_bit_f32(%qh1, %qh_bit) : (i8, i32) -> (f32) + %high2 = func.call @ggml_q5k_split_wmma_high_bit_f32(%qh2, %qh_bit) : (i8, i32) -> (f32) + %high3 = func.call @ggml_q5k_split_wmma_high_bit_f32(%qh3, %qh_bit) : (i8, i32) -> (f32) + %q0_low = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1_low = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2_low = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3_low = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %q0 = scalar.addf %q0_low, %high0 : f32 + %q1 = scalar.addf %q1_low, %high1 : f32 + %q2 = scalar.addf %q2_low, %high2 : f32 + %q3 = scalar.addf %q3_low, %high3 : f32 + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %half0 = scalar.fptrunc %value0 : f32 to f16 + %half1 = scalar.fptrunc %value1 : f32 to f16 + %half2 = scalar.fptrunc %value2 : f32 to f16 + %half3 = scalar.fptrunc %value3 : f32 to f16 + %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> + func.return %result : vector<4xf16> +} + +// Acquires one naturally aligned Q5_K block header as a single 16-byte packet. +func.def inline @ggml_q5k_split_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q5_block: index) -> (vector<4xi32>) { + %c0 = index.constant 0 : index + %block_bytes = index.constant 176 : offset + %block_byte_add = index.scale %q5_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> + %header_words = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> + func.return %header_words : vector<4xi32> +} + +template.decl @ggml.mul_mat_id_q5k_f16_wmma.project_vector16(%stage_activation: i1, %q5_group: index, %k_origin: index, %load_row: index, %load_k: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %header0: vector<4xi32>, %header1: vector<4xi32>, %q_word0: vector<1xi32>, %q_word1: vector<1xi32>, %qh_bytes0: vector<4xi8>, %qh_bytes1: vector<4xi8>, %accumulator: vector<16xf16>) -> (vector<16xf16>) + +template.def<@ggml.mul_mat_id_q5k_f16_wmma.project_vector16> device @ggml_mul_mat_id_q5k_f16_wmma_project_vector16(%stage_activation: i1, %q5_group: index, %k_origin: index, %load_row: index, %load_k: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %header0: vector<4xi32>, %header1: vector<4xi32>, %q_word0: vector<1xi32>, %q_word1: vector<1xi32>, %qh_bytes0: vector<4xi8>, %qh_bytes1: vector<4xi8>, %accumulator: vector<16xf16>) -> (vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %input_noalias, %route_stage_noalias, %weight_stage_noalias, %activation_stage_noalias = buffer.assume.noalias %input, %route_stage, %weight_stage, %activation_stage : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> + %route_stage_view = buffer.view %route_stage_noalias[%c0_offset] : buffer -> view<32xi32> + %weight_stage_view = buffer.view %weight_stage_noalias[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %bounded_group = index.assume %q5_group [range(%q5_group, 0, 7)] : 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 + %selects_first_row = index.cmp eq, %row_offset, %c0 : index + %selected_header = scf.select %selects_first_row, %header0, %header1 : vector<4xi32> + %selected_q_word = scf.select %selects_first_row, %q_word0, %q_word1 : vector<1xi32> + %selected_qh_bytes = scf.select %selects_first_row, %qh_bytes0, %qh_bytes1 : vector<4xi8> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %decoded = func.call @ggml_q5k_split_wmma_vector4_from_header_code(%bounded_group, %selected_header, %selected_q_word, %selected_qh_bytes) : (index, vector<4xi32>, vector<1xi32>, vector<4xi8>) -> (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> + } + scf.if %stage_activation { + %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %route_count : index + %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next_accumulator = scf.for %k_half = [%c0 to %c32 step %c16](%half_accumulator = %accumulator : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %next = vector.mma %lhs, %rhs, %half_accumulator : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return %next_accumulator : vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q6k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q6k.loom new file mode 100644 index 000000000000..022a7ab1b63c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_projection_q6k.loom @@ -0,0 +1,169 @@ +// Shared routed F16 pair Q6_K single-projection WMMA step. +// +// Callers own the outer routing and fusion schedule. This motif stages one +// projected Q6_K weight tile, optionally stages the shared routed activation +// tile, performs the 16x16 WMMA update, and returns the updated accumulator. + +// Acquires the packed words shared by one four-group half of a Q6_K block. +// Each QL word supplies two groups and the QH word supplies all four. +func.def inline @ggml_q6k_split_wmma_load_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_half: index, %packet: index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %bounded_half = index.assume %q6_half [range(%q6_half, 0, 1)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %ql_half_word_base = index.mul %bounded_half, %c16 : index + %ql0_word_index = index.add %ql_half_word_base, %bounded_packet : index + %ql1_word_base = index.add %ql_half_word_base, %c8 : index + %ql1_word_index = index.add %ql1_word_base, %bounded_packet : index + %qh_half_word_base = index.mul %bounded_half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %ql0_word = vector.load %ql_view[%ql0_word_index] : view<32xi32> -> vector<1xi32> + %ql1_word = vector.load %ql_view[%ql1_word_index] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + func.return %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> +} + +// Decodes four adjacent values after the surrounding schedule has selected +// the scale and retained the packed code words at their natural lifetimes. +func.def inline @ggml_q6k_split_wmma_vector4_from_scale_codes(%q6_group: index, %scale_i8: i8, %d_f16: f16, %ql_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %group_in_half = index.rem %bounded_group, %c4 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +// Loads only the group scale and block multiplier while reusing packed words +// retained by an enclosing four-group half-block schedule. +func.def inline @ggml_q6k_split_wmma_vector4_from_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index, %ql0_word: vector<1xi32>, %ql1_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 210 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %uses_ql1 = index.cmp eq, %ql_side, %c1 : index + %ql_word = scf.select %uses_ql1, %ql1_word, %ql0_word : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index = index.add %scale_group_base, %scale_packet_half : index + %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %values = func.call @ggml_q6k_split_wmma_vector4_from_scale_codes(%bounded_group, %scale_i8, %d_f16, %ql_word, %qh_word) : (index, i8, f16, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) + func.return %values : vector<4xf16> +} + +template.decl @ggml.mul_mat_id_q6k_f16_wmma.project_vector16(%stage_activation: i1, %q6_group: index, %q6_block: index, %k_origin: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %ql0_word0: vector<1xi32>, %ql1_word0: vector<1xi32>, %qh_word0: vector<1xi32>, %ql0_word1: vector<1xi32>, %ql1_word1: vector<1xi32>, %qh_word1: vector<1xi32>, %accumulator: vector<16xf16>) -> (vector<16xf16>) + +template.def<@ggml.mul_mat_id_q6k_f16_wmma.project_vector16> device @ggml_mul_mat_id_q6k_f16_wmma_project_vector16(%stage_activation: i1, %q6_group: index, %q6_block: index, %k_origin: index, %load_row: index, %load_k: index, %load_packet: index, %channel_tile_base: index, %bounded_output_size: index, %assignment_count: index, %token_count: index, %route_count: index, %input_size: index, %subgroup_channel_add: index, %subgroup_route_add: index, %input: buffer, %route_stage: buffer, %weight_stage: buffer, %activation_stage: buffer, %weight: buffer, %expert_byte_base: offset, %weight_row_bytes: offset, %ql0_word0: vector<1xi32>, %ql1_word0: vector<1xi32>, %qh_word0: vector<1xi32>, %ql0_word1: vector<1xi32>, %ql1_word1: vector<1xi32>, %qh_word1: vector<1xi32>, %accumulator: vector<16xf16>) -> (vector<16xf16>) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_i32 = scalar.constant 0 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %input_noalias, %route_stage_noalias, %weight_stage_noalias, %activation_stage_noalias, %weight_noalias = buffer.assume.noalias %input, %route_stage, %weight_stage, %activation_stage, %weight : buffer, buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> + %route_stage_view = buffer.view %route_stage_noalias[%c0_offset] : buffer -> view<32xi32> + %weight_stage_view = buffer.view %weight_stage_noalias[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage_noalias[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : 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 + %selects_first_row = index.cmp eq, %row_offset, %c0 : index + %selected_ql0_word = scf.select %selects_first_row, %ql0_word0, %ql0_word1 : vector<1xi32> + %selected_ql1_word = scf.select %selects_first_row, %ql1_word0, %ql1_word1 : vector<1xi32> + %selected_qh_word = scf.select %selects_first_row, %qh_word0, %qh_word1 : vector<1xi32> + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %decoded = func.call @ggml_q6k_split_wmma_vector4_from_half_codes(%weight_noalias, %row_byte_base, %q6_block, %bounded_group, %load_packet, %selected_ql0_word, %selected_ql1_word, %selected_qh_word) : (buffer, offset, index, index, index, vector<1xi32>, vector<1xi32>, vector<1xi32>) -> (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> + } + scf.if %stage_activation { + %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %route_count : index + %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next_accumulator = scf.for %k_half = [%c0 to %c32 step %c16](%half_accumulator = %accumulator : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %next = vector.mma %lhs, %rhs, %half_accumulator : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.return %next_accumulator : vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_publish.loom new file mode 100644 index 000000000000..fad4270b6c50 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/f16_publish.loom @@ -0,0 +1,25 @@ +// F16 publication for routed down-projection results. +template.decl @ggml.mul_mat_id_f16_f16_wmma.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %channel_tile: index, %expert: index, %route_tile_base: index, %expert_route_count: index, %expert_table: buffer, %output: buffer) + +template.decl @ggml.mul_mat_id_f16_f16_wmma.publish_vector4(%publish_word: i1, %assignment: index, %channel: index, %token_count: index, %route_count: index, %output_size: index, %values: vector<4xf16>, %output: buffer) + +template.def<@ggml.mul_mat_id_f16_f16_wmma.publish_vector4> device @ggml_mul_mat_id_f16_f16_wmma_publish_vector4(%publish_word: i1, %assignment: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf16>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 8)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 4096)] : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %assignment_count = index.mul %token_count, %route_count : index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> + scf.if %publish_word { + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<4xf16>, view<[%assignment_count]x[%output_size]xf16>, vector<4xi1> + } + template.return +} + +template.def<@ggml.mul_mat_id_f16_f16_wmma.finish_tile> device @ggml_mul_mat_id_f16_f16_wmma_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %channel_tile: index, %expert: index, %route_tile_base: index, %expert_route_count: index, %expert_table: buffer, %output: buffer) { + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..11f2ebbcf68a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/pair_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,341 @@ +// Shared routed F32 gate/up schedule and dual accumulation. +template.decl @ggml.matmul_id.pair_f32_swiglu.apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.pair_f32_swiglu.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c9_i32 = scalar.constant 9 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c511_i32 = scalar.constant 511 : i32 + %expert_i32 = scalar.andi %descriptor, %c511_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c9_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c15_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 511)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +template.def<@ggml.matmul_id.pair_f32_swiglu.body> device @ggml_mul_mat_id_pair_f32_swiglu_f32_f32_wmma_body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 32)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index + %bounded_input_route_count = index.assume %input_route_count [range(%input_route_count, 1, 32), le(%input_route_count, %bounded_route_count)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %output_row_count = index.mul %bounded_token_count, %bounded_route_count : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %gate_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%gate_weight_format, %bounded_input_size) : (index, index) -> (offset) + %up_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%up_weight_format, %bounded_input_size) : (index, index) -> (offset) + %gate_weight_expert_bytes = index.scale %bounded_output_size, %gate_weight_row_bytes : index, offset -> offset + %up_weight_expert_bytes = index.scale %bounded_output_size, %up_weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %partition_table_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_route_count]x[%bounded_input_size]xf32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%bounded_token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf32> + %gate_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %up_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %gate_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %up_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %gate_weight_stage_view = buffer.view %gate_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %up_weight_stage_view = buffer.view %up_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %gate_result_fragment_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %up_result_fragment_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %gate_result_physical_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %up_result_physical_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [le(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @ggml_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %bounded_token_count [lt(%assignment_ordinal, %bounded_token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%bounded_token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %gate_expert_byte_base = index.scale %bounded_expert, %gate_weight_expert_bytes : index, offset -> offset + %up_expert_byte_base = index.scale %bounded_expert, %up_weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %gate_result00, %gate_result01, %gate_result10, %gate_result11, %up_result00, %up_result01, %up_result10, %up_result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%gate_block_acc00 = %init00 : vector<8xf16>, %gate_block_acc01 = %init01 : vector<8xf16>, %gate_block_acc10 = %init10 : vector<8xf16>, %gate_block_acc11 = %init11 : vector<8xf16>, %up_block_acc00 = %init00 : vector<8xf16>, %up_block_acc01 = %init01 : vector<8xf16>, %up_block_acc10 = %init10 : vector<8xf16>, %up_block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%gate_acc00 = %gate_block_acc00 : vector<8xf16>, %gate_acc01 = %gate_block_acc01 : vector<8xf16>, %gate_acc10 = %gate_block_acc10 : vector<8xf16>, %gate_acc11 = %gate_block_acc11 : vector<8xf16>, %up_acc00 = %up_block_acc00 : vector<8xf16>, %up_acc01 = %up_block_acc01 : vector<8xf16>, %up_acc10 = %up_block_acc10 : vector<8xf16>, %up_acc11 = %up_block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %gate_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %gate_weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %gate_expert_byte_base, %channel_byte_add : offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %up_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %up_weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %up_expert_byte_base, %channel_byte_add : offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %gate_weight_values, %gate_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + vector.store %up_weight_values, %up_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_assignment, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %route0 = index.rem %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_route0 = index.rem %route0, %bounded_input_route_count : index + %input_route = index.assume %input_route0 [range(%input_route0, 0, 31), lt(%input_route0, %bounded_input_route_count)] : index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %loaded = vector.load.mask %input_view[%token, %input_route, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_route_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 = scf.for %k_half = [%c0 to %c32 step %c16](%gate_half_acc00 = %gate_acc00 : vector<8xf16>, %gate_half_acc01 = %gate_acc01 : vector<8xf16>, %gate_half_acc10 = %gate_acc10 : vector<8xf16>, %gate_half_acc11 = %gate_acc11 : vector<8xf16>, %up_half_acc00 = %up_acc00 : vector<8xf16>, %up_half_acc01 = %up_acc01 : vector<8xf16>, %up_half_acc10 = %up_acc10 : vector<8xf16>, %up_half_acc11 = %up_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %gate_lhs0 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %gate_lhs1 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs0 = vector.fragment.load %up_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs1 = vector.fragment.load %up_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %gate_half_next00 = vector.mma %gate_lhs0, %rhs0, %gate_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next01 = vector.mma %gate_lhs0, %rhs1, %gate_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next10 = vector.mma %gate_lhs1, %rhs0, %gate_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next11 = vector.mma %gate_lhs1, %rhs1, %gate_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next00 = vector.mma %up_lhs0, %rhs0, %up_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next01 = vector.mma %up_lhs0, %rhs1, %up_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next10 = vector.mma %up_lhs1, %rhs0, %up_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next11 = vector.mma %up_lhs1, %rhs1, %up_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %gate_half_next00, %gate_half_next01, %gate_half_next10, %gate_half_next11, %up_half_next00, %up_half_next01, %up_half_next10, %up_half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { + scf.yield %assignment0_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { + scf.yield %assignment1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 65535)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 65535)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %gate_result00, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result00, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result01, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result01, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result10, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result10, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result11, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result11, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %gate_values = vector.load %gate_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.matmul_id.pair_f32_swiglu.apply_vector4>(%gate_wide, %up_wide) : (vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/pair_swiglu_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/pair_swiglu_publish_f32.loom new file mode 100644 index 000000000000..9a150fe939d6 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/pair_swiglu_publish_f32.loom @@ -0,0 +1,14 @@ +// Shared SwiGLU activation for routed F32 pair publication. +template.decl @ggml.matmul_id.pair_f32_swiglu.apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) + +template.def<@ggml.matmul_id.pair_f32_swiglu.apply_vector4> device @ggml_mul_mat_id_pair_f32_swiglu_apply_vector4(%gate: vector<4xf32>, %up: vector<4xf32>) -> (vector<4xf32>) { + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c1_f32x4 = vector.constant 1.0 : vector<4xf32> + %negative_gate = vector.subf %c0_f32x4, %gate : vector<4xf32> + %gate_exp = vector.expf %negative_gate : vector<4xf32> + %denominator = vector.addf %c1_f32x4, %gate_exp : vector<4xf32> + %sigmoid = vector.divf %c1_f32x4, %denominator : vector<4xf32> + %silu = vector.mulf %gate, %sigmoid : vector<4xf32> + %result = vector.mulf %silu, %up : vector<4xf32> + template.return %result : vector<4xf32> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..2c59b2d7d059 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,302 @@ +// Shared routed expert WMMA matmul pipeline for GGML MUL_MAT_ID. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul_id.skinny_f32_f32.publish_vector4 +// ggml.matmul_id.skinny_f32_f32.finish_tile +template.decl @ggml.matmul_id.skinny_f32_f32.core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +// Common descriptor packing uses a 9-bit expert ordinal, a 6-bit 32-row +// partition ordinal, and a 5-bit row count minus one. +func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c9_i32 = scalar.constant 9 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c511_i32 = scalar.constant 511 : i32 + %expert_i32 = scalar.andi %descriptor, %c511_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c9_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c15_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 511)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +template.def<@ggml.matmul_id.skinny_f32_f32.core> device @ggml_mul_mat_id_skinny_f32_f32_core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index + %bounded_input_route_count = index.assume %input_route_count [range(%input_route_count, 1, 32), le(%input_route_count, %bounded_route_count)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %weight_block_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) + %quant_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %quant_block_count, %weight_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %partition_table_noalias, %weight_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %weight : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%bounded_token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [le(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @ggml_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %bounded_token_count [lt(%assignment_ordinal, %bounded_token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%bounded_token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %weight_k = index.add %k_origin, %load_k : index + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %route0 = index.rem %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_route0 = index.rem %route0, %bounded_input_route_count : index + %input_route = index.assume %input_route0 [range(%input_route0, 0, 31), lt(%input_route0, %bounded_input_route_count)] : index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_route, %input_k] : view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { + scf.yield %assignment0_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { + scf.yield %assignment1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 65535)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 65535)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %assignment_token0 = index.div %bounded_assignment0, %bounded_route_count : index + %assignment_route0 = index.rem %bounded_assignment0, %bounded_route_count : index + %assignment_token1 = index.div %bounded_assignment1, %bounded_route_count : index + %assignment_route1 = index.rem %bounded_assignment1, %bounded_route_count : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes00, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes01, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes10, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.skinny_f32_f32.publish_vector4>(%writes11, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.matmul_id.skinny_f32_f32.finish_tile>(%bounded_token_count, %bounded_route_count, %bounded_output_size, %bounded_expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %channel_tile, %descriptor_ordinal, %bounded_expert, %route_tile_base, %partition_row_count, %expert_table, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom new file mode 100644 index 000000000000..e07d95cd30e1 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/skinny_postprocess_f32.loom @@ -0,0 +1,167 @@ +template.decl @ggml.matmul_id.skinny_f32_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul_id.skinny_f32_f32.add_bias_vector4> device @ggml_mul_mat_id_skinny_f32_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul_id.skinny_f32_f32.add_residual_vector4> device @ggml_mul_mat_id_skinny_f32_f32_add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %row_count = index.mul %token_count, %route_count : index + %assignment, %assignment_count = index.assume %assignment0, %row_count [lt(%assignment0, %row_count)] : index, index + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%assignment, %channel], %mask, %c0_f32x4 : view<[%row_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul_id.skinny_f32_f32.store_vector4> device @ggml_mul_mat_id_skinny_f32_f32_store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %row_count = index.mul %token_count, %route_count : index + %assignment, %assignment_count = index.assume %assignment0, %row_count [lt(%assignment0, %row_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%assignment, %channel], %mask : vector<4xf32>, view<[%row_count]x[%output_size]xf32>, vector<4xi1> + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight> device @ggml_mul_mat_id_skinny_f32_f32_finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index + %expert_count = index.assume %expert_count0 [range(%expert_count0, 1, 512)] : index + %output_tile_count = index.assume %output_tile_count0 [range(%output_tile_count0, 1, 512)] : index + %maximum_partition_count = index.assume %maximum_partition_count0 [range(%maximum_partition_count0, 1, 2560)] : index + %descriptor_ordinal, %table_partition_count = index.assume %descriptor_ordinal0, %maximum_partition_count [lt(%descriptor_ordinal0, %maximum_partition_count)] : index, index + %expert, %table_expert_count = index.assume %expert0, %expert_count [lt(%expert0, %expert_count)] : index, index + %route_tile_base, %table_token_count0 = index.assume %route_tile_base0, %token_count [lt(%route_tile_base0, %token_count)] : index, index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %counter_scratch_bytes = index.constant 4 : offset + %rms_scratch_bytes = index.constant 512 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %assignment_table_byte_base = index.scale %expert_count, %c4_bytes : index, offset -> offset + %row_count = index.mul %token_count, %route_count : index + %hidden_size_i32 = index.cast %output_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_channel_ordinal_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %expert_table_noalias, %source_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counters_noalias = buffer.assume.noalias %expert_table, %source, %norm_weight, %normalized_output, %completion_counters : buffer, buffer, buffer, buffer, buffer + %completion_counters_aligned = buffer.assume.alignment %completion_counters_noalias {minimum_alignment = 16} : buffer + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%expert_count]x[%token_count]xi32> + %source_view = buffer.view %source_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %norm_weight_view = buffer.view %norm_weight_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %normalized_output_view = buffer.view %normalized_output_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %completion_counter_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%maximum_partition_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %rms_scratch = buffer.alloca align(16) %rms_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %rms_scratch_view = buffer.view %rms_scratch[%c0_offset] : buffer -> view<128xf32> + kernel.barrier scope(workgroup) ordering(release) + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%descriptor_ordinal] {ordering = acq_rel, scope = device} : i32, view<[%maximum_partition_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %is_last_channel_tile = scalar.cmpi eq, %old_counter, %last_channel_ordinal_i32 : i32 + scf.if %is_last_channel_tile { + kernel.barrier scope(workgroup) ordering(acquire) + scf.for %row_offset = [%c0 to %c32 step %c1] { + %valid_row = index.cmp ult, %row_offset, %partition_row_count : index + scf.if %valid_row { + %assignment_ordinal0 = index.add %route_tile_base, %row_offset : index + %assignment_ordinal, %table_token_count = index.assume %assignment_ordinal0, %token_count [lt(%assignment_ordinal0, %token_count)] : index, index + %assignment_i32 = view.load %assignment_view[%expert, %assignment_ordinal] : view<[%expert_count]x[%token_count]xi32> -> i32 + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %assignment_count = index.assume %assignment, %row_count [lt(%assignment, %row_count)] : index, index + %thread_sum = scf.for %channel = [%workitem to %output_size step %c128](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %source_view[%bounded_assignment, %channel] : view<[%row_count]x[%output_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %rms_scratch_view[%subgroup] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c2 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %rms_scratch_view[%lane] : view<128xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %rms_scratch_view[%c0] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %rms_scratch_view[%c0] : view<128xf32> -> f32 + scf.for %channel = [%workitem to %output_size step %c128] { + %value = view.load %source_view[%bounded_assignment, %channel] : view<[%row_count]x[%output_size]xf32> -> f32 + %weight_value = view.load %norm_weight_view[%channel] : view<[%output_size]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %normalized = scalar.mulf %scaled, %weight_value : f32 + view.store %normalized, %normalized_output_view[%bounded_assignment, %channel] : f32, view<[%row_count]x[%output_size]xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%descriptor_ordinal] {ordering = release, scope = device} : i32, view<[%maximum_partition_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom new file mode 100644 index 000000000000..359fcf33e146 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_input_f32_dequant_accumulate_f32.loom @@ -0,0 +1,302 @@ +// Shared routed expert WMMA matmul pipeline for GGML MUL_MAT_ID. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.matmul_id.tiled_f32_f32.publish_vector4 +// ggml.matmul_id.tiled_f32_f32.finish_tile +template.decl @ggml.matmul_id.tiled_f32_f32.core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: vector<4xf32>, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +// Common descriptor packing uses a 9-bit expert ordinal, a 6-bit 32-row +// partition ordinal, and a 5-bit row count minus one. +func.def inline @ggml_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { + %c1_i32 = scalar.constant 1 : i32 + %c5_i32 = scalar.constant 5 : i32 + %c9_i32 = scalar.constant 9 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c31_i32 = scalar.constant 31 : i32 + %c63_i32 = scalar.constant 63 : i32 + %c511_i32 = scalar.constant 511 : i32 + %expert_i32 = scalar.andi %descriptor, %c511_i32 : i32 + %partition_shifted_i32 = scalar.shrui %descriptor, %c9_i32 : i32 + %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 + %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 + %row_count_shifted_i32 = scalar.shrui %descriptor, %c15_i32 : i32 + %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 + %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 511)] : index + %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index + %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index + %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + func.return %expert, %route_tile_base, %partition_row_count : index, index, index +} + +template.def<@ggml.matmul_id.tiled_f32_f32.core> device @ggml_mul_mat_id_tiled_f32_f32_core(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %epsilon: f32, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 32)] : index + %bounded_input_route_count = index.assume %input_route_count [range(%input_route_count, 1, 32), le(%input_route_count, %bounded_route_count)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %route_stage_bytes = index.constant 128 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_i32 = scalar.constant 0 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index + %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset + %weight_block_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) + %quant_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.scale %quant_block_count, %weight_block_bytes : index, offset -> offset + %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset + %input_noalias, %expert_table_noalias, %partition_table_noalias, %weight_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %weight : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%bounded_token_count]xi32> + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %is_workitem_zero = index.cmp eq, %workitem, %c0 : index + %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 + %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [le(%partition_count0, %maximum_partition_count)] : index, index + scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { + %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { + %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 + %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 + %expert, %route_tile_base, %partition_row_count = func.call @ggml_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) + %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index + %loads_route = index.cmp ult, %workitem, %c32 : index + scf.if %loads_route { + %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index + %assignment_ordinal = index.add %route_tile_base, %local_route : index + %valid_row = index.cmp ult, %local_route, %partition_row_count : index + %assignment_i32 = scf.if %valid_row -> (i32) { + %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %bounded_token_count [lt(%assignment_ordinal, %bounded_token_count)] : index, index + %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%bounded_token_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values = scf.if %valid_channel -> (vector<4xf16>) { + %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset + %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset + %weight_k = index.add %k_origin, %load_k : index + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %is_activation_row = index.cmp ult, %local_row, %c32 : index + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 + %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 + %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index + %token0 = index.div %bounded_assignment, %bounded_route_count : index + %route0 = index.rem %bounded_assignment, %bounded_route_count : index + %token, %input_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_route0 = index.rem %route0, %bounded_input_route_count : index + %input_route = index.assume %input_route0 [range(%input_route0, 0, 31), lt(%input_route0, %bounded_input_route_count)] : index + %input_k = index.add %k_origin, %load_k : index + %loaded = vector.load %input_view[%token, %input_route, %input_k] : view<[%bounded_token_count]x[%bounded_input_route_count]x[%input_size]xf32> -> vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_route0 = index.div %lane, %c4 : index + %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %local_route1 = index.add %c16, %publish_route : index + %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 + %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 + %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 + %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 + %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { + scf.yield %assignment0_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { + scf.yield %assignment1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index + %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index + %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 65535)] : index + %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 65535)] : index + %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index + %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index + %assignment_token0 = index.div %bounded_assignment0, %bounded_route_count : index + %assignment_route0 = index.rem %bounded_assignment0, %bounded_route_count : index + %assignment_token1 = index.div %bounded_assignment1, %bounded_route_count : index + %assignment_route1 = index.rem %bounded_assignment1, %bounded_route_count : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 + %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 + %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 + %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 + vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes00, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes01, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel0, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes10, %bounded_assignment0, %assignment_token0, %assignment_route0, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %wide = vector.extf %values : vector<4xf16> to vector<4xf32> + template.apply<@ggml.matmul_id.tiled_f32_f32.publish_vector4>(%writes11, %bounded_assignment1, %assignment_token1, %assignment_route1, %channel1, %bounded_token_count, %bounded_route_count, %bounded_output_size, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.matmul_id.tiled_f32_f32.finish_tile>(%bounded_token_count, %bounded_route_count, %bounded_output_size, %bounded_expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %channel_tile, %descriptor_ordinal, %bounded_expert, %route_tile_base, %partition_row_count, %expert_table, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_postprocess_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_postprocess_f32.loom new file mode 100644 index 000000000000..3514bac420a7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/tiled_postprocess_f32.loom @@ -0,0 +1,167 @@ +template.decl @ggml.matmul_id.tiled_f32_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.tiled_f32_f32.add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.matmul_id.tiled_f32_f32.add_bias_vector4> device @ggml_mul_mat_id_tiled_f32_f32_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul_id.tiled_f32_f32.add_residual_vector4> device @ggml_mul_mat_id_tiled_f32_f32_add_residual_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %row_count = index.mul %token_count, %route_count : index + %assignment, %assignment_count = index.assume %assignment0, %row_count [lt(%assignment0, %row_count)] : index, index + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%assignment, %channel], %mask, %c0_f32x4 : view<[%row_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.matmul_id.tiled_f32_f32.store_vector4> device @ggml_mul_mat_id_tiled_f32_f32_store_vector4(%token_count0: index, %route_count0: index, %output_size0: index, %assignment0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %row_count = index.mul %token_count, %route_count : index + %assignment, %assignment_count = index.assume %assignment0, %row_count [lt(%assignment0, %row_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%assignment, %channel], %mask : vector<4xf32>, view<[%row_count]x[%output_size]xf32>, vector<4xi1> + template.return +} + +template.def<@ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight> device @ggml_mul_mat_id_tiled_f32_f32_finish_routed_rmsnorm_weight(%token_count0: index, %route_count0: index, %output_size0: index, %expert_count0: index, %output_tile_count0: index, %maximum_partition_count0: index, %epsilon: f32, %descriptor_ordinal0: index, %expert0: index, %route_tile_base0: index, %partition_row_count0: index, %expert_table: buffer, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index + %expert_count = index.assume %expert_count0 [range(%expert_count0, 1, 512)] : index + %output_tile_count = index.assume %output_tile_count0 [range(%output_tile_count0, 1, 512)] : index + %maximum_partition_count = index.assume %maximum_partition_count0 [range(%maximum_partition_count0, 1, 2560)] : index + %descriptor_ordinal, %table_partition_count = index.assume %descriptor_ordinal0, %maximum_partition_count [lt(%descriptor_ordinal0, %maximum_partition_count)] : index, index + %expert, %table_expert_count = index.assume %expert0, %expert_count [lt(%expert0, %expert_count)] : index, index + %route_tile_base, %table_token_count0 = index.assume %route_tile_base0, %token_count [lt(%route_tile_base0, %token_count)] : index, index + %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %counter_scratch_bytes = index.constant 4 : offset + %rms_scratch_bytes = index.constant 512 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %assignment_table_byte_base = index.scale %expert_count, %c4_bytes : index, offset -> offset + %row_count = index.mul %token_count, %route_count : index + %hidden_size_i32 = index.cast %output_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_channel_ordinal_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %expert_table_noalias, %source_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counters_noalias = buffer.assume.noalias %expert_table, %source, %norm_weight, %normalized_output, %completion_counters : buffer, buffer, buffer, buffer, buffer + %completion_counters_aligned = buffer.assume.alignment %completion_counters_noalias {minimum_alignment = 16} : buffer + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%expert_count]x[%token_count]xi32> + %source_view = buffer.view %source_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %norm_weight_view = buffer.view %norm_weight_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %normalized_output_view = buffer.view %normalized_output_noalias[%c0_offset] : buffer -> view<[%row_count]x[%output_size]xf32> + %completion_counter_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%maximum_partition_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %rms_scratch = buffer.alloca align(16) %rms_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %rms_scratch_view = buffer.view %rms_scratch[%c0_offset] : buffer -> view<128xf32> + kernel.barrier scope(workgroup) ordering(release) + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%descriptor_ordinal] {ordering = acq_rel, scope = device} : i32, view<[%maximum_partition_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %is_last_channel_tile = scalar.cmpi eq, %old_counter, %last_channel_ordinal_i32 : i32 + scf.if %is_last_channel_tile { + kernel.barrier scope(workgroup) ordering(acquire) + scf.for %row_offset = [%c0 to %c32 step %c1] { + %valid_row = index.cmp ult, %row_offset, %partition_row_count : index + scf.if %valid_row { + %assignment_ordinal0 = index.add %route_tile_base, %row_offset : index + %assignment_ordinal, %table_token_count = index.assume %assignment_ordinal0, %token_count [lt(%assignment_ordinal0, %token_count)] : index, index + %assignment_i32 = view.load %assignment_view[%expert, %assignment_ordinal] : view<[%expert_count]x[%token_count]xi32> -> i32 + %assignment0 = index.cast %assignment_i32 : i32 to index + %assignment = index.assume %assignment0 [range(%assignment0, 0, 65535)] : index + %bounded_assignment, %assignment_count = index.assume %assignment, %row_count [lt(%assignment, %row_count)] : index, index + %thread_sum = scf.for %channel = [%workitem to %output_size step %c128](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %source_view[%bounded_assignment, %channel] : view<[%row_count]x[%output_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %rms_scratch_view[%subgroup] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c2 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %rms_scratch_view[%lane] : view<128xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %rms_scratch_view[%c0] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %rms_scratch_view[%c0] : view<128xf32> -> f32 + scf.for %channel = [%workitem to %output_size step %c128] { + %value = view.load %source_view[%bounded_assignment, %channel] : view<[%row_count]x[%output_size]xf32> -> f32 + %weight_value = view.load %norm_weight_view[%channel] : view<[%output_size]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %normalized = scalar.mulf %scaled, %weight_value : f32 + view.store %normalized, %normalized_output_view[%bounded_assignment, %channel] : f32, view<[%row_count]x[%output_size]xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%descriptor_ordinal] {ordering = release, scope = device} : i32, view<[%maximum_partition_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_pair.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_pair.loom new file mode 100644 index 000000000000..ebfb6aae8b82 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_pair.loom @@ -0,0 +1,279 @@ +// Common token-vector routed pair contraction and binary publication. +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +template.decl @ggml.matmul_id.vector_pair.body(%publish_output: i1, %binary_op: index, %input_size: index, %token_count: index, %token: index, %route_count: index, %route: index, %route_stride: index, %expert_count: index, %output_size: index, %channel: index, %lane: index, %q8_input: buffer, %route_ids: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) + +template.decl @ggml.matmul_id.vector_pair.binary(%binary_op: index, %lhs: f32, %rhs: f32) -> (f32) + +template.def<@ggml.matmul_id.vector_pair.binary> device @ggml_matmul_id_vector_pair_binary(%binary_op: index, %lhs: f32, %rhs: f32) -> (f32) { + %swiglu_op = index.constant 4 : index + %is_swiglu = index.cmp eq, %binary_op, %swiglu_op : index + %result = scf.if %is_swiglu -> (f32) { + // Preserve the established routed SwiGLU lowering for numerical and HSACO parity. + %lhs_silu = scalar.siluf %lhs : f32 + %swiglu = scalar.mulf %lhs_silu, %rhs : f32 + scf.yield %swiglu : f32 + } else { + %generic = template.apply<@ggml.binary_f32.apply>(%binary_op, %lhs, %rhs) : (index, f32, f32) -> (f32) + scf.yield %generic : f32 + } + template.return %result : f32 +} + +template.def<@ggml.matmul_id.vector_pair.body> device @ggml_matmul_id_vector_pair_body(%publish_output: i1, %binary_op: index, %input_size: index, %token_count: index, %token: index, %route_count: index, %route: index, %route_stride: index, %expert_count: index, %output_size: index, %channel: index, %lane: index, %q8_input: buffer, %route_ids: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512)] : index + %bounded_route_count, %bounded_route_stride = index.assume %route_count, %route_stride [range(%route_count, 1, 8), range(%route_stride, 1, 128), le(%route_count, %route_stride)] : index, index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, 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 + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %q4_block_bytes = index.constant 144 : index + %q4_block_bytes_offset = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %c1_byte = index.constant 1 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %q4_block_count = index.div %input_size, %c256 : index + %weight_row_bytes = index.mul %q4_block_count, %q4_block_bytes : index + %weight_expert_bytes = index.mul %bounded_output_size, %weight_row_bytes : index + %q8_group_count = index.div %input_size, %c128 : index + %q8_bytes_per_token = index.mul %q8_group_count, %q4_block_bytes : index + %padded_input_size = index.add %input_size, %c1023 : index + %iteration_count = index.div %padded_input_size, %c1024 : index + %q8_noalias, %route_ids_noalias, %lhs_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %lhs_weight, %rhs_weight, %output : buffer, buffer, buffer, buffer, buffer + %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_count]x[%bounded_output_size]xf32> + %route_index, %route_row_stride = index.assume %bounded_route, %bounded_route_stride [lt(%bounded_route, %bounded_route_stride)] : index, index + %expert_i32 = view.load %route_ids_view[%bounded_token, %route_index] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index + %expert_byte_base = index.mul %expert, %weight_expert_bytes : index + %channel_byte_add = index.mul %bounded_channel, %weight_row_bytes : index + %row_byte_index = index.add %expert_byte_base, %channel_byte_add : index + %row_byte_base = index.scale %row_byte_index, %c1_byte : index, offset -> offset + %q8_token_byte_index = index.mul %bounded_token, %q8_bytes_per_token : index + %q8_token_byte_base = index.scale %q8_token_byte_index, %c1_byte : index, offset -> offset + %q4_group_pair0 = index.div %bounded_lane, %c2 : index + %q4_group_pair1 = index.rem %q4_group_pair0, %c4 : index + %q4_group_pair = index.assume %q4_group_pair1 [range(%q4_group_pair1, 0, 3)] : index + %q4_half0 = index.rem %bounded_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %lane_q4_block = index.div %bounded_lane, %c8 : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %gate_acc, %up_acc = scf.for %iteration = [%c0 to %iteration_count step %c1](%gate_iter = %c0_f32 : f32, %up_iter = %c0_f32 : f32) -> (f32, f32) unroll { + %iteration_q4_block = index.mul %iteration, %c4 : index + %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index + %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index + %gate_contribution, %up_contribution = scf.if %valid_q4_block -> (f32, f32) { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %q8_block_group_base = index.mul %q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_token_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_noalias[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_noalias[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %q4_block, %q4_block_bytes_offset : index, offset -> offset + %q4_block_byte_base = index.add %row_byte_base, %q4_block_byte_add : offset + %gate_header_view = buffer.view %lhs_noalias[%q4_block_byte_base] : buffer -> view<4xi32> + %up_header_view = buffer.view %rhs_noalias[%q4_block_byte_base] : buffer -> view<4xi32> + %gate_header_words = vector.load %gate_header_view[0] : view<4xi32> -> vector<4xi32> + %up_header_words = vector.load %up_header_view[0] : view<4xi32> -> vector<4xi32> + %gate_q4_low, %gate_low_d_scale, %gate_low_dmin_scale, %gate_q4_high, %gate_high_d_scale, %gate_high_dmin_scale = func.call @ggml_matmul_id_q4k_chunk_pair(%lhs_noalias, %row_byte_base, %q4_block, %q4_group_pair, %q4_half, %gate_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %up_q4_low, %up_low_d_scale, %up_low_dmin_scale, %up_q4_high, %up_high_d_scale, %up_high_dmin_scale = func.call @ggml_matmul_id_q4k_chunk_pair(%rhs_noalias, %row_byte_base, %q4_block, %q4_group_pair, %q4_half, %up_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %gate_low = func.call @ggml_matmul_id_q4k_q8_1_dot(%gate_q4_low, %gate_low_d_scale, %gate_low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %gate_high = func.call @ggml_matmul_id_q4k_q8_1_dot(%gate_q4_high, %gate_high_d_scale, %gate_high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %up_low = func.call @ggml_matmul_id_q4k_q8_1_dot(%up_q4_low, %up_low_d_scale, %up_low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %up_high = func.call @ggml_matmul_id_q4k_q8_1_dot(%up_q4_high, %up_high_d_scale, %up_high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %gate_pair = scalar.addf %gate_low, %gate_high : f32 + %up_pair = scalar.addf %up_low, %up_high : f32 + scf.yield %gate_pair, %up_pair : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + %gate_next = scalar.addf %gate_iter, %gate_contribution : f32 + %up_next = scalar.addf %up_iter, %up_contribution : f32 + scf.yield %gate_next, %up_next : f32, f32 + } + %gate_dot = kernel.subgroup.reduce %gate_acc : f32 + %up_dot = kernel.subgroup.reduce %up_acc : f32 + %lane_i32 = index.cast %bounded_lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %writes_output = scalar.andi %publish_output, %is_lane_zero : i1 + scf.if %writes_output { + %result = template.apply<@ggml.matmul_id.vector_pair.binary>(%binary_op, %gate_dot, %up_dot) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%bounded_token, %bounded_route, %bounded_channel] : f32, view<[%bounded_token_count]x[%bounded_route_count]x[%bounded_output_size]xf32> + } + template.return +} + +// Strict Q8_1 x4 publication for routed vector pair producers. +template.decl @ggml.matmul_id.vector_pair.q8_1_x4.subgroup_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) + +template.decl @ggml.matmul_id.vector_pair.q8_1_x4.completion(%token_count: index, %route_count: index, %output_size: index, %token: index, %route: index, %channel_base: index, %subgroup: index, %workitem: index, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) + +// One selected subgroup packs a 128-value physical group as four Q8_1 blocks. +template.def<@ggml.matmul_id.vector_pair.q8_1_x4.subgroup_body> device @ggml_matmul_id_vector_pair_q8_1_x4_subgroup_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 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 + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127_f32 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %group_element_base = index.mul %group, %c128 : index + %lane_element_add = index.mul %lane, %c4 : index + %input_index = index.add %group_element_base, %lane_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_maximum = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %block0 = index.div %lane, %c8 : index + %block = index.assume %block0 [range(%block0, 0, 3)] : index + %block_maxima = vector.insert %thread_maximum into %c0_f32x4[%block] : f32, vector<4xf32> + %reduced_maxima = kernel.subgroup.reduce %block_maxima : vector<4xf32> + %amax = vector.extract %reduced_maxima[%block] : vector<4xf32> -> f32 + %d = scalar.divf %amax, %c127_f32 : f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + scf.if %publish_output { + vector.store %packed_word, %group_qs[%lane] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + %block_sums0 = vector.insert %thread_sum into %c0_f32x4[%block] : f32, vector<4xf32> + %reduced_sums = kernel.subgroup.reduce %block_sums0 : vector<4xf32> + %quantized_sum = vector.extract %reduced_sums[%block] : vector<4xf32> -> f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %word_in_block = index.rem %lane, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %publishes_metadata = scalar.andi %publish_output, %is_block_leader : i1 + scf.if %publishes_metadata { + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + template.return +} + +// Every producer workgroup advances one four-channel arrival. The final +// arrival acquires the full physical group, publishes Q8, and resets it. +template.def<@ggml.matmul_id.vector_pair.q8_1_x4.completion> device @ggml_matmul_id_vector_pair_q8_1_x4_completion(%token_count: index, %route_count: index, %output_size: index, %token: index, %route: index, %channel_base: index, %subgroup: index, %workitem: index, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c0_i32 = scalar.constant 0 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %physical_group_count = index.div %output_size, %c128 : index + %row_count = index.mul %token_count, %route_count : index + %completion_counter_count = index.mul %row_count, %physical_group_count : index + %token_row_base = index.mul %token, %route_count : index + %row = index.add %token_row_base, %route : index + %row_group_base = index.mul %row, %physical_group_count : index + %group_in_row = index.div %channel_base, %c128 : index + %counter_index0 = index.add %row_group_base, %group_in_row : index + %counter_index, %bounded_completion_counter_count = index.assume %counter_index0, %completion_counter_count [lt(%counter_index0, %completion_counter_count)] : index, index + %completion_counters_aligned = buffer.assume.alignment %completion_counters {minimum_alignment = 16} : buffer + %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%bounded_completion_counter_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + %old_counter0 = view.atomic.rmw %c4_i32, %completion_counters_view[%counter_index] {ordering = acq_rel, scope = device} : i32, view<[%bounded_completion_counter_count]xi32> -> i32 + view.store %old_counter0, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %group_size_i32 = index.cast %c128 : index to i32 + %last_arrival_i32 = scalar.subi %group_size_i32, %c4_i32 : i32 + %negative_group_size_i32 = scalar.subi %c0_i32, %group_size_i32 : i32 + %is_last_arrival = scalar.cmpi eq, %old_counter, %last_arrival_i32 : i32 + scf.if %is_last_arrival { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_output = scalar.constant true : i1 + %is_quantize_subgroup = index.cmp eq, %subgroup, %c0 : index + scf.if %is_quantize_subgroup { + template.apply<@ggml.matmul_id.vector_pair.q8_1_x4.subgroup_body>(%publish_output, %bounded_completion_counter_count, %counter_index, %output, %next_q8_output) : (i1, index, index, buffer, buffer) + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_group_size_i32, %completion_counters_view[%counter_index] {ordering = release, scope = device} : i32, view<[%bounded_completion_counter_count]xi32> + } + } + template.return +} +func.decl @ggml_matmul_id_q4k_chunk_pair(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) +func.decl @ggml_matmul_id_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_q4k_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_q4k_q8.loom new file mode 100644 index 000000000000..d150a0950c25 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_q4k_q8.loom @@ -0,0 +1,183 @@ +// Q4_K decode and Q8_1 dot-product primitives for routed vector matmul. +func.def inline @ggml_matmul_id_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +func.def inline @ggml_matmul_id_q4k_chunk_pair(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %code_byte_add = index.constant 16 : offset + %c4_i32 = scalar.constant 4 : i32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index + %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_byte_add : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> + %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 + %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 + %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 + %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 + %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %pair_code_base = index.mul %bounded_pair, %c8 : index + %half_code_add = index.mul %bounded_half, %c4 : index + %code_index0 = index.add %pair_code_base, %half_code_add : index + %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index + %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> + %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> + %c4_i32v = vector.splat %c4_i32 : vector<4xi32> + %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> + %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> + %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> + %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> + %low_group = index.mul %bounded_pair, %c2 : index + %high_group = index.add %low_group, %c1 : index + %low_scale, %low_minimum = func.call @ggml_matmul_id_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) + %high_scale, %high_minimum = func.call @ggml_matmul_id_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) + %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 + %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 + %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 + %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 + %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 + %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 + %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 + %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 + func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 +} + +func.def inline @ggml_matmul_id_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { + %c0_i32 = scalar.constant 0 : i32 + %c0_i32v = vector.constant 0 : vector<4xi32> + %half_f32 = scalar.constant 0.5 : f32 + %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> + %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 + %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 + %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 + %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 + %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 + %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 + %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 + func.return %contribution : f32 +} + +// Computes two adjacent Q4_K group contributions for one lane in a block. +func.def inline @ggml_matmul_id_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index + %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : 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 + %c256 = index.constant 256 : index + %q4_block_bytes = index.constant 144 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_byte_add = index.constant 16 : offset + %q4_block_count = index.div %bounded_input_size, %c256 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index + %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index + %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index + %q4_half0 = index.rem %bounded_block_lane, %c2 : index + %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index + %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index + %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index + %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index + %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index + %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index + %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index + %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index + %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index + %q8_half_word_add = index.mul %q4_half, %c4 : index + %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index + %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index + %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index + %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index + %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index + %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index + %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index + %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index + %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index + %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index + %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset + %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> + %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 + %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 + %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 + %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 + %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 + %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 + %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 + %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 + %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> + %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> + %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> + %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> + %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset + %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset + %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> + %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> + %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @ggml_matmul_id_q4k_chunk_pair(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) + %low = func.call @ggml_matmul_id_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %high = func.call @ggml_matmul_id_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) + %pair = scalar.addf %low, %high : f32 + func.return %pair : f32 +} + +// An eight-lane cohort contracts one complete routed row. +func.def inline @ggml_matmul_id_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) { + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_cohort_lane = index.assume %cohort_lane [range(%cohort_lane, 0, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %sum = scf.for %q4_block0 = [%c0 to %q4_block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) unroll { + %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index + %pair = func.call @ggml_matmul_id_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %bounded_cohort_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) + %next = scalar.addf %block_acc, %pair : f32 + scf.yield %next : f32 + } + func.return %sum : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_weighted.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_weighted.loom new file mode 100644 index 000000000000..5e49241dc4a1 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul_id/vector_weighted.loom @@ -0,0 +1,258 @@ +// Routed down projection with route weighting, residual add, and completion. +template.decl @ggml.matmul_id.vector_weighted.q4k_q8.wave32_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) + +func.decl @ggml_matmul_id_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) + +template.def<@ggml.matmul_id.vector_weighted.q4k_q8.wave32_body> device @ggml_matmul_id_vector_weighted_q4k_q8_wave32_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index + %channel_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %c144 = index.constant 144 : index + %c256 = index.constant 256 : index + %c1_byte = index.constant 1 : offset + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %cohort0 = index.div %lane, %c8 : index + %cohort = index.assume %cohort0 [range(%cohort0, 0, 3)] : index + %cohort_lane0 = index.rem %lane, %c8 : index + %cohort_lane = index.assume %cohort_lane0 [range(%cohort_lane0, 0, 7)] : index + %channel_base = index.mul %channel_tile, %c8 : index + %channel = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %lane_i32 = index.cast %lane : index to i32 + %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 + %q4_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_byte_count = index.mul %q4_block_count, %c144 : index + %q8_group_count = index.div %bounded_input_size, %c128 : index + %q8_row_byte_count = index.mul %q8_group_count, %c144 : index + %padded_route_count = index.add %bounded_route_count, %c3 : index + %route_batch_count = index.div %padded_route_count, %c4 : index + %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer + %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> + %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> + %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index + %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index + %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 + %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 + %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { + %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index + %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index + %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 + %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 + scf.yield %expert_i32, %route_weight : i32, f32 + } else { + scf.yield %c0_i32, %c0_f32 : i32, f32 + } + %computes_channel = scalar.andi %publish_output, %valid_channel : i1 + %routed_lane_sum = scf.if %computes_channel -> (f32) { + %sum = scf.for %route_batch = [%c0 to %route_batch_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { + %route_batch_base = index.mul %route_batch, %c4 : index + %route0 = index.add %route_batch_base, %cohort : index + %route = index.assume %route0 [range(%route0, 0, 7)] : index + %active_route = index.cmp ult, %route, %bounded_route_count : index + %weighted_lane = scf.if %active_route -> (f32) { + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index + %route_i32 = index.cast %bounded_route : index to i32 + %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index + %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index + %expert_output_base = index.mul %expert, %bounded_output_size : index + %expert_channel = index.add %expert_output_base, %channel : index + %weight_row_byte_index = index.mul %expert_channel, %weight_row_byte_count : index + %weight_row_byte_base = index.scale %weight_row_byte_index, %c1_byte : index, offset -> offset + %q8_row_base0 = index.mul %bounded_token, %body_route_count : index + %q8_row = index.add %q8_row_base0, %bounded_route : index + %q8_row_byte_index = index.mul %q8_row, %q8_row_byte_count : index + %q8_row_byte_base = index.scale %q8_row_byte_index, %c1_byte : index, offset -> offset + %route_lane_sum = func.call @ggml_matmul_id_q4k_q8_1_x4_cohort_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %cohort_lane) : (index, buffer, offset, buffer, offset, index) -> (f32) + %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 + %weighted = scalar.mulf %route_lane_sum, %route_weight : f32 + scf.yield %weighted : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next = scalar.addf %route_acc, %weighted_lane : f32 + scf.yield %next : f32 + } + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 + %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 + %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 + scf.if %writes_output { + %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index + %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 + %result = scalar.addf %residual, %routed_sum : f32 + view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> + } + template.return +} + +// Wave64 Q6_K by F32 routed down projection with paired channel publication. +template.decl @ggml.matmul_id.vector_weighted.q6k_f32.wave64_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) + +func.decl @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) + +func.decl @ggml_q6k_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>) + +func.decl @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) + +template.def<@ggml.matmul_id.vector_weighted.q6k_f32.wave64_pair_body> device @ggml_matmul_id_vector_weighted_q6k_f32_wave64_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 2, 4096)] : index + %pair_tile = kernel.workgroup.id : index + %subgroup0 = kernel.subgroup.id : index + %lane0 = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %scale_stage_bytes = index.constant 2048 : offset + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %pair_tile_base = index.mul %pair_tile, %c4 : index + %pair = index.add %pair_tile_base, %subgroup : index + %row00 = index.mul %pair, %c2 : index + %row0, %body_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %body_output_size : index + %safe_row1 = scf.select %row1_valid, %row1, %c0 : index + %assignment_count = index.mul %body_token_count, %bounded_route_count : index + %q6_block_count = index.div %bounded_input_size, %c256 : index + %cohort = index.div %lane, %c16 : index + %scale_frame0 = index.mul %subgroup, %c2 : index + %scale_frame1 = index.add %scale_frame0, %c1 : index + %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer + %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> + %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%body_output_size]xf32> + %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer + %routed_sum0, %routed_sum1 = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc0 = %c0_f32 : f32, %route_acc1 = %c0_f32 : f32) -> (f32, f32) { + %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index + %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index + %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 + %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 + %expert0 = index.cast %expert_i32 : i32 to index + %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index + %expert, %body_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index + %expert_output_base = index.mul %expert, %body_output_size : index + %expert_row0 = index.add %expert_output_base, %row0 : index + %expert_row1 = index.add %expert_output_base, %safe_row1 : index + %input_row_base = index.mul %bounded_token, %body_route_count : index + %input_row = index.add %input_row_base, %bounded_route : index + %lane_sum0, %lane_sum1 = scf.for %block_base = [%c0 to %q6_block_count step %c4](%block_acc0 = %c0_f32 : f32, %block_acc1 = %c0_f32 : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + %input00, %input01, %input02, %input03 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage, %input00, %input01, %input02, %input03) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) + // Reload after the first scale stage so the first row's vectors die. + %input10, %input11, %input12, %input13 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %row1_contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage, %input10, %input11, %input12, %input13) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %c0_f32 : f32 + } + %next0 = scalar.addf %block_acc0, %contribution0 : f32 + %next1 = scalar.addf %block_acc1, %contribution1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %route_sum0 = kernel.subgroup.reduce %lane_sum0 : f32 + %route_sum1 = kernel.subgroup.reduce %lane_sum1 : f32 + %weighted0 = scalar.mulf %route_sum0, %route_weight : f32 + %weighted1 = scalar.mulf %route_sum1, %route_weight : f32 + %next0 = scalar.addf %route_acc0, %weighted0 : f32 + %next1 = scalar.addf %route_acc1, %weighted1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + %residual0 = view.load %output_view[%bounded_token, %row0] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 + %result0 = scalar.addf %residual0, %routed_sum0 : f32 + view.store %result0, %output_view[%bounded_token, %row0] : f32, view<[%body_token_count]x[%body_output_size]xf32> + scf.if %row1_valid { + %residual1 = view.load %output_view[%bounded_token, %row1] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 + %result1 = scalar.addf %residual1, %routed_sum1 : f32 + view.store %result1, %output_view[%bounded_token, %row1] : f32, view<[%body_token_count]x[%body_output_size]xf32> + } + } + template.return +} + +// Completion protocol shared by routed down projections publishing next Q8. +template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) + +template.decl @ggml.matmul_id.vector_weighted.next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) + +template.def<@ggml.matmul_id.vector_weighted.next_q8_completion> device @ggml_matmul_id_vector_weighted_next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %output_channels_per_workgroup = index.assume %output_channels_per_workgroup0 [range(%output_channels_per_workgroup0, 1, 8)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 128, 32768), mul(%output_size, 128)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %c0 = index.constant 0 : index + %workitem = kernel.workitem.id : index + %output_tile_count = index.div %bounded_output_size, %output_channels_per_workgroup : index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %output_noalias, %norm_weight_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter0 = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter0, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 + scf.if %is_last_output_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_normalized = scalar.constant false : i1 + template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %reduction_subgroup_count, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %next_q8_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_postops.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_postops.loom new file mode 100644 index 000000000000..ade604180adb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_postops.loom @@ -0,0 +1,149 @@ +template.decl @ggml.mul_mat_f32_f32_wmma.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.mul_mat_f32_f32_wmma.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.mul_mat_f32_f32_wmma.finish_rmsnorm_weight(%token_count0: index, %output_size0: index, %output_tile_count0: index, %token_tile_count0: index, %epsilon: f32, %token_tile: index, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.mul_mat_f32_f32_wmma.add_bias_vector4> device @ggml_mul_mat_f32_f32_wmma_add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) { + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %bias_noalias = buffer.assume.noalias %bias : buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %bias_values = vector.load.mask %bias_view[%channel], %mask, %c0_f32x4 : view<[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %bias_values : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.mul_mat_f32_f32_wmma.add_residual_vector4> device @ggml_mul_mat_f32_f32_wmma_add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %values, %residual : vector<4xf32> + template.return %sum : vector<4xf32> +} + +template.def<@ggml.mul_mat_f32_f32_wmma.store_vector4> device @ggml_mul_mat_f32_f32_wmma_store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 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 + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.return +} + +template.def<@ggml.mul_mat_f32_f32_wmma.finish_rmsnorm_weight> device @ggml_mul_mat_f32_f32_wmma_finish_rmsnorm_weight(%token_count0: index, %output_size0: index, %output_tile_count0: index, %token_tile_count0: index, %epsilon: f32, %token_tile: index, %source: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index + %output_tile_count = index.assume %output_tile_count0 [range(%output_tile_count0, 1, 4096)] : index + %token_tile_count = index.assume %token_tile_count0 [range(%token_tile_count0, 1, 64)] : index + %workitem = kernel.workitem.id : index + %workgroup_size = kernel.workgroup.size : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %subgroup_count = index.div %workgroup_size, %c64 : index + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %rms_scratch_bytes = index.constant 512 : offset + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %hidden_size_i32 = index.cast %output_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_channel_ordinal_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %source_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counters_noalias = buffer.assume.noalias %source, %norm_weight, %normalized_output, %completion_counters : buffer, buffer, buffer, buffer + %completion_counters_aligned = buffer.assume.alignment %completion_counters_noalias {minimum_alignment = 16} : buffer + %source_view = buffer.view %source_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %norm_weight_view = buffer.view %norm_weight_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %normalized_output_view = buffer.view %normalized_output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %completion_counter_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%token_tile_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %rms_scratch = buffer.alloca align(16) %rms_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + %rms_scratch_view = buffer.view %rms_scratch[%c0_offset] : buffer -> view<128xf32> + %token_tile_base = index.mul %token_tile, %c32 : index + kernel.barrier scope(workgroup) ordering(release) + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + scf.if %workitem_is_zero { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%token_tile] {ordering = acq_rel, scope = device} : i32, view<[%token_tile_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %is_last_channel_tile = scalar.cmpi eq, %old_counter, %last_channel_ordinal_i32 : i32 + scf.if %is_last_channel_tile { + kernel.barrier scope(workgroup) ordering(acquire) + scf.for %token_offset = [%c0 to %c32 step %c1] { + %token0 = index.add %token_tile_base, %token_offset : index + %valid_token = index.cmp ult, %token0, %token_count : index + scf.if %valid_token { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %thread_sum = scf.for %channel = [%workitem to %output_size step %workgroup_size](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %source_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %rms_scratch_view[%subgroup] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %rms_scratch_view[%lane] : view<128xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %rms_scratch_view[%c0] : f32, view<128xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %rms_scratch_view[%c0] : view<128xf32> -> f32 + scf.for %channel = [%workitem to %output_size step %workgroup_size] { + %value = view.load %source_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %weight_value = view.load %norm_weight_view[%channel] : view<[%output_size]xf32> -> f32 + %scaled = scalar.mulf %value, %scale : f32 + %normalized = scalar.mulf %scaled, %weight_value : f32 + view.store %normalized, %normalized_output_view[%token, %channel] : f32, view<[%token_count]x[%output_size]xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %workitem_is_zero { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%token_tile] {ordering = release, scope = device} : i32, view<[%token_tile_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom new file mode 100644 index 000000000000..f96948802fd4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_f32_f32_wmma_core.loom @@ -0,0 +1,1611 @@ +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 WMMA matmul pipeline. +// +// Op wrappers provide publication and tile-finalization hooks by defining: +// ggml.mul_mat_f32_f32_wmma.publish_vector4 +// ggml.mul_mat_f32_f32_wmma.finish_tile +template.decl @ggml.mul_mat_f32_f32_wmma.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: 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.mul_mat_f32_f32_wmma.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.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: vector<4xf32>, %arg8: buffer, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: 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_mul_mat_f32_f32_wmma_tiled(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_output_size = index.add %bounded_output_size, %c63 : index + %output_tile_count = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %bounded_token_count, %c31 : index + %token_tile_count = index.div %padded_token_count, %c32 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %bounded_input_size) : (index, index) -> (offset) + %split_count = kernel.workgroup.count : index + %partition = kernel.workgroup.id : index + %blocks_per_partition = index.div %quant_block_count, %split_count : index + %quant_block_begin = index.mul %partition, %blocks_per_partition : index + %quant_block_end = index.add %quant_block_begin, %blocks_per_partition : index + %input_noalias, %weight_noalias = buffer.assume.noalias %input, %weight : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_token, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %loaded = vector.load.mask %input_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %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.mul_mat_f32_f32_wmma.publish_vector4>(%writes00, %token0, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + 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.mul_mat_f32_f32_wmma.publish_vector4>(%writes01, %token1, %channel0, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + 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.mul_mat_f32_f32_wmma.publish_vector4>(%writes10, %token0, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + 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.mul_mat_f32_f32_wmma.publish_vector4>(%writes11, %token1, %channel1, %bounded_token_count, %bounded_output_size, %output_accumulation, %output_unary_op, %wide, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + template.apply<@ggml.mul_mat_f32_f32_wmma.finish_tile>(%bounded_token_count, %bounded_output_size, %output_tile_count, %token_tile_count, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def pure inline @ggml_q8_lowtoken_geometry(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %q4_format = index.constant 44 : index + %q6_format = index.constant 46 : index + %q8_format = index.constant 9 : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %native_q4 = index.cmp eq, %weight_format, %q4_format : index + %native_q6 = index.cmp eq, %weight_format, %q6_format : index + %q8_input = index.cmp eq, %activation_format, %q8_format : index + %q4_min_blocks = index.constant 327680 : index + %blocks = index.div %input_size, %c256 : index + %weight_blocks = index.mul %blocks, %output_size : index + %large_payload = index.cmp uge, %weight_blocks, %q4_min_blocks : index + %q4_can_stage = scalar.andi %native_q4, %large_payload : i1 + %native_staged = scalar.ori %q4_can_stage, %native_q6 : i1 + %native_q8 = scalar.andi %native_staged, %q8_input : i1 + %q4_min_k = index.constant 16384 : index + %q6_min_k = index.constant 8192 : index + %min_k = scf.select %native_q4, %q4_min_k, %q6_min_k : index + %min_n = index.constant 4096 : index + %five_tokens = index.cmp eq, %token_capacity, %c5 : index + %long_k = index.cmp uge, %input_size, %min_k : index + %wide_n = index.cmp uge, %output_size, %min_n : index + %contracts = index.cmp ule, %output_size, %input_size : index + %large = scalar.andi %long_k, %wide_n : i1 + %long_reduction = scalar.andi %large, %contracts : i1 + %staged_shape = scalar.andi %five_tokens, %long_reduction : i1 + %staged = scalar.andi %native_q8, %staged_shape : i1 + %rows = scf.select %native_q4, %c64, %c32 : index + %result = scf.select %staged, %rows, %c0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %block_bytes = index.constant 144 : index + %half_cache_bytes = index.constant 16777216 : index + %cache_bytes = index.constant 33554432 : index + %payload_bytes = index.mul %weight_blocks, %block_bytes : index + %large_half = index.cmp uge, %payload_bytes, %half_cache_bytes : index + %fits_cache = index.cmp ule, %payload_bytes, %cache_bytes : index + %cache_range = scalar.andi %large_half, %fits_cache : i1 + %remainder = index.rem %blocks, %c2 : index + %even_blocks = index.cmp eq, %remainder, %c0 : index + %reduces_rows = index.cmp ult, %output_size, %input_size : index + %native_q4_q8 = scalar.andi %native_q4, %q8_input : i1 + %plain = index.cmp eq, %result, %c0 : index + %q4_plain = scalar.andi %native_q4_q8, %plain : i1 + %five_wide = scalar.andi %five_tokens, %wide_n : i1 + %reduction = scalar.andi %reduces_rows, %even_blocks : i1 + %partition_shape = scalar.andi %five_wide, %reduction : i1 + %partition_payload = scalar.andi %partition_shape, %cache_range : i1 + %partitioned = scalar.andi %q4_plain, %partition_payload : i1 + %k_partitions = scf.select %partitioned, %c2, %c1 : index + func.return %result, %k_partitions : index, index +} + +func.def pure inline @ggml_mul_mat_uses_lowtoken_dot(%token_capacity: index, %input_size: index) -> (i1) { + %c0 = index.constant 0 : index + %c5 = index.constant 5 : index + %c256 = index.constant 256 : index + %few_tokens = index.cmp ule, %token_capacity, %c5 : index + %tail = index.rem %input_size, %c256 : index + %whole_blocks = index.cmp eq, %tail, %c0 : index + %use_dot = scalar.andi %few_tokens, %whole_blocks : i1 + func.return %use_dot : i1 +} + +template.def<@ggml.mul_mat_f32_f32_wmma.launch> pure @ggml_mul_mat_f32_f32_wmma_launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, 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 + %c5 = index.constant 5 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %lowtoken = func.call pure @ggml_mul_mat_uses_lowtoken_dot(%token_capacity, %input_size) : (index, index) -> (i1) + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %small_output = index.cmp ult, %output_size, %c64 : index + %narrow = scalar.andi %lowtoken, %small_output : i1 + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%token_capacity, %input_size, %output_size, %weight_format) : (index, index, index, index) -> (index, index) + %staged = index.cmp ugt, %staged_rows, %c0 : index + %staged_rounding = index.sub %staged_rows, %c1 : index + %staged_threads = index.mul %staged_rows, %c16 : index + %cohort_rows = index.div %c8, %k_partitions : index + %cohort_rounding = index.sub %cohort_rows, %c1 : index + %base_row_tile = scf.select %narrow, %c4, %cohort_rows : index + %c1024 = index.constant 1024 : index + %c4096 = index.constant 4096 : index + %activation_format = config.get @ggml.mul_mat.activation_format : index + %f32_input = index.cmp eq, %activation_format, %c0 : index + %batched_tokens = index.cmp ugt, %token_capacity, %c1 : index + %long_reduction = index.cmp uge, %input_size, %c4096 : index + %partition_tail = index.rem %input_size, %c1024 : index + %full_partitions = index.cmp eq, %partition_tail, %c0 : index + %batch_f32 = scalar.andi %batched_tokens, %f32_input : i1 + %partitionable = scalar.andi %long_reduction, %full_partitions : i1 + %partition_rows = scalar.andi %batch_f32, %partitionable : i1 + %narrow_threads = scf.select %partition_rows, %c1024, %c256 : index + %base_workgroup_size = scf.select %narrow, %narrow_threads, %c128 : index + %base_rounding = scf.select %narrow, %c3, %cohort_rounding : index + %row_rounding = scf.select %staged, %staged_rounding, %base_rounding : index + %row_tile = scf.select %staged, %staged_rows, %base_row_tile : index + %workgroup_size = scf.select %staged, %staged_threads, %base_workgroup_size : index + %rounded_n = index.add %output_size, %row_rounding : index + %direct_tiles = index.div %rounded_n, %row_tile : index + %launch_x = scf.select %lowtoken, %direct_tiles, %output_tiles : index + %launch_y = scf.select %lowtoken, %c1, %token_tiles : index + template.return %launch_x, %launch_y, %c1, %workgroup_size : index, index, index, index +} + +func.def inline @ggml_lowtoken_combine_k_partitions(%value: f32, %partitions: index) -> (f32) { + %c1 = index.constant 1 : index + %partitioned = index.cmp ugt, %partitions, %c1 : index + %result = scf.if %partitioned -> (f32) { + %offset = scalar.constant 16 : i32 + %width = scalar.constant 32 : i32 + %peer, %ok = kernel.subgroup.shuffle %value, %offset, %width : f32, i32, i32 + %sum = scalar.addf %value, %peer : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_lowtoken_reduce_cohort_f32(%value: f32) -> (f32) { + %width = scalar.constant 32 : i32 + %off1 = scalar.constant 1 : i32 + %peer1, %ok1 = kernel.subgroup.shuffle %value, %off1, %width : f32, i32, i32 + %sum1 = scalar.addf %value, %peer1 : f32 + %off2 = scalar.constant 2 : i32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %off2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %off4 = scalar.constant 4 : i32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %off4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %off8 = scalar.constant 8 : i32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %off8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + func.return %sum8 : f32 +} + +// Narrow low-token projections assign a full wave to each row. +// Both cohort sizes use the existing wrapper publication hooks. +func.def inline @ggml_lowtoken_reduce_f32(%values: vector<4xf32>, %full_wave: i1) -> (f32) { + %zero = scalar.constant 0.0 : f32 + %width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %sum = vector.reduce %values, %zero : vector<4xf32>, f32 + %peer1, %ok1 = kernel.subgroup.shuffle %sum, %xor1, %width : f32, i32, i32 + %sum1 = scalar.addf %sum, %peer1 : f32 + %peer2, %ok2 = kernel.subgroup.shuffle %sum1, %xor2, %width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %peer2 : f32 + %peer4, %ok4 = kernel.subgroup.shuffle %sum2, %xor4, %width : f32, i32, i32 + %sum4 = scalar.addf %sum2, %peer4 : f32 + %peer8, %ok8 = kernel.subgroup.shuffle %sum4, %xor8, %width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer8 : f32 + %result = scf.if %full_wave -> (f32) { + %xor16 = scalar.constant 16 : i32 + %peer16, %ok16 = kernel.subgroup.shuffle %sum8, %xor16, %width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer16 : f32 + %lane32 = scalar.constant 32 : i32 + %peer32 = kernel.subgroup.broadcast %sum16 from %lane32 : f32, i32 + %sum64 = scalar.addf %sum16, %peer32 : f32 + scf.yield %sum64 : f32 + } else { + scf.yield %sum8 : f32 + } + func.return %result : f32 +} + +func.def inline @ggml_mul_mat_f32_f32_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %tokens = index.assume %token_count [range(%token_count, 1, 5)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workgroup_size = kernel.workgroup.size : index + %one_wave_per_row = index.cmp eq, %workgroup_size, %c256 : index + %c1024 = index.constant 1024 : index + %partition_group = index.cmp eq, %workgroup_size, %c1024 : index + %narrow_output = index.cmp ult, %n, %c64 : index + %partitioned = scalar.andi %partition_group, %narrow_output : i1 + %full_wave = scalar.ori %one_wave_per_row, %partitioned : i1 + %k_partitions = scf.select %partitioned, %c4, %c1 : index + %cohort_width = scf.select %full_wave, %c64, %c16 : index + %tile_rows = scf.select %full_wave, %c4, %c8 : index + %wave_rows = scf.select %full_wave, %c1, %c4 : index + %step_width = scf.select %full_wave, %c256, %c64 : index + %steps_per_block = scf.select %full_wave, %c1, %c4 : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %lane16 = index.rem %lane, %cohort_width : index + %cohort = index.div %lane, %cohort_width : index + %tile_base = index.mul %channel_tile, %tile_rows : index + %row_subgroup = index.div %subgroup, %k_partitions : index + %k_partition = index.rem %subgroup, %k_partitions : index + %wave_add = index.mul %row_subgroup, %wave_rows : index + %wave_base = index.add %tile_base, %wave_add : index + %row = index.add %wave_base, %cohort : index + %valid_row = index.cmp ult, %row, %n : index + %packet = index.rem %lane16, %c8 : index + %lane_group = index.div %lane16, %c8 : index + %lane_k = index.mul %lane16, %c4 : index + %blocks = index.div %k, %c256 : index + %steps = index.div %k, %step_width : index + %partition_steps = index.div %steps, %k_partitions : index + %step_begin = index.mul %k_partition, %partition_steps : index + %step_end = index.add %step_begin, %partition_steps : index + %weight_bytes = func.call @ggml_dequant_weight_tile_bytes(%weight_format) : (index) -> (offset) + %row_bytes = index.scale %blocks, %weight_bytes : index, offset -> offset + %row_base = index.scale %row, %row_bytes : index, offset -> offset + %a_na, %w_na = buffer.assume.noalias %input, %weight : buffer, buffer + %av = buffer.view %a_na[%base] : buffer -> view<[%tokens]x[%k]xf32> + %total0, %total1, %total2, %total3, %total4 = scf.for %step = [%step_begin to %step_end step %c1](%acc0 = %zero : vector<4xf32>, %acc1 = %zero : vector<4xf32>, %acc2 = %zero : vector<4xf32>, %acc3 = %zero : vector<4xf32>, %acc4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %block = index.div %step, %steps_per_block : index + %quarter = index.rem %step, %steps_per_block : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %step_k = index.mul %step, %step_width : index + %kk = index.add %step_k, %lane_k : index + %ww = scf.if %valid_row -> (vector<4xf32>) { + %format44 = index.constant 44 : index + %packed = index.cmp eq, %weight_format, %format44 : index + %wh = scf.if %packed -> (vector<4xf32>) { + %v = func.call @ggml_q4k_native_row64_f32_vector4(%w_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } else { + %v = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %w_na, %row_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %v : vector<4xf32> + } + scf.yield %wh : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next0 = scf.if %valid0 -> (vector<4xf32>) { + %a = vector.load %av[%t0, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc0 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc0 : vector<4xf32> + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next1 = scf.if %valid1 -> (vector<4xf32>) { + %a = vector.load %av[%t1, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc1 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc1 : vector<4xf32> + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next2 = scf.if %valid2 -> (vector<4xf32>) { + %a = vector.load %av[%t2, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc2 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc2 : vector<4xf32> + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next3 = scf.if %valid3 -> (vector<4xf32>) { + %a = vector.load %av[%t3, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc3 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc3 : vector<4xf32> + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next4 = scf.if %valid4 -> (vector<4xf32>) { + %a = vector.load %av[%t4, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %r = vector.fmaf %ww, %a, %acc4 : vector<4xf32> + scf.yield %r : vector<4xf32> + } else { + scf.yield %acc4 : vector<4xf32> + } + scf.yield %next0, %next1, %next2, %next3, %next4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %leader = index.cmp eq, %lane16, %c0 : index + %c5 = index.constant 5 : index + %stage_width = scf.select %partitioned, %c16, %c8 : index + %stage_elements = index.mul %stage_width, %c5 : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%stage_width]xf32> + %local_row0 = index.add %wave_add, %cohort : index + %partial_offset = index.mul %k_partition, %c4 : index + %local_row = index.add %local_row0, %partial_offset : index + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total0, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total1, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total2, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total3, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum8 = func.call @ggml_lowtoken_reduce_f32(%total4, %full_wave) : (vector<4xf32>, i1) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%stage_width]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %tile_rows, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %tile_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, 4), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %stage_width : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_origin = scf.if %partitioned -> (index) { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 64), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } else { + %origin = index.assume %stage_origin0 [range(%stage_origin0, 0, 36), mul(%stage_origin0, 4)] : index + scf.yield %origin : index + } + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values0 = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + %values = scf.if %partitioned -> (vector<4xf32>) { + %stage_origin1 = index.add %stage_origin, %c4 : index + %stage_origin2 = index.add %stage_origin1, %c4 : index + %stage_origin3 = index.add %stage_origin2, %c4 : index + %values1 = vector.load %flat_stage[%stage_origin1] : view<[%stage_elements]xf32> -> vector<4xf32> + %values2 = vector.load %flat_stage[%stage_origin2] : view<[%stage_elements]xf32> -> vector<4xf32> + %values3 = vector.load %flat_stage[%stage_origin3] : view<[%stage_elements]xf32> -> vector<4xf32> + %sum01 = vector.addf %values0, %values1 : vector<4xf32> + %sum012 = vector.addf %sum01, %values2 : vector<4xf32> + %sum = vector.addf %sum012, %values3 : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values0 : vector<4xf32> + } + template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %row_rounding = index.sub %tile_rows, %c1 : index + %rounded_n = index.add %n, %row_rounding : index + %output_tiles = index.div %rounded_n, %tile_rows : index + template.apply<@ggml.mul_mat_f32_f32_wmma.finish_tile>(%tokens, %n, %output_tiles, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Copy the original Q8_1_x4 records for up to 2048 reduction elements. +// Call uniformly at a slab boundary; only the final slab may be partial. +func.def inline @ggml_q8_1_x4_stage_lowtoken_slab(%input: buffer, %stage: buffer, %capacity: index, %slab_begin: index, %blocks: index, %global_row_bytes: offset) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c144 = index.constant 144 : index + %c576 = index.constant 576 : index + %c18 = index.constant 18 : index + %base = index.constant 0 : offset + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %workgroup_size = kernel.workgroup.size : index + %thread = kernel.workitem.id : index + %stage_packets = index.mul %capacity, %c144 : index + %stage_elements = index.mul %capacity, %c576 : index + %stage_view = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xi32> + %slab_limit = index.add %slab_begin, %c8 : index + %full_slab = index.cmp ule, %slab_limit, %blocks : index + %slab_end = scf.select %full_slab, %slab_limit, %blocks : index + %active_blocks0 = index.sub %slab_end, %slab_begin : index + %active_blocks = index.assume %active_blocks0 [range(%active_blocks0, 1, 8)] : index + %active_packets = index.mul %active_blocks, %c18 : index + %slab_byte_add = index.scale %slab_begin, %block_bytes : index, offset -> offset + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %copy_packet = [%thread to %stage_packets step %workgroup_size] { + %copy_token = index.div %copy_packet, %c144 : index + %copy_inner = index.rem %copy_packet, %c144 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + scf.if %copy_valid { + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %copy_bytes = index.scale %copy_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slab_byte_add : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + %copy_origin = index.mul %copy_packet, %c4 : index + vector.store %values, %stage_view[%copy_origin] : vector<4xi32>, view<[%stage_elements]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_mul_mat_q8_lowtoken_publish(%tokens: index, %n: index, %row_tile: index, %row_base: index, %local_row: index, %block_lane: index, %channel_tile: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %total_g0: f32, %total_g1: f32, %total_g2: f32, %total_g3: f32, %total_g4: f32, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %leader = index.cmp eq, %block_lane, %c0 : index + %width = scalar.constant 32 : i32 + %stage_rows = index.constant 5 : index + %stage_elements = index.mul %stage_rows, %row_tile : index + %element_bytes = index.constant 4 : offset + %stage_bytes = index.scale %stage_elements, %element_bytes : index, offset -> offset + %stage = buffer.alloca align(16) %stage_bytes : buffer + %sv = buffer.view %stage[%base] : buffer -> view<5x[%row_tile]xf32> + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %sum = scalar.addf %total_g0, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p0, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %sum = scalar.addf %total_g1, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p1, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %sum = scalar.addf %total_g2, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p2, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %sum = scalar.addf %total_g3, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p3, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %sum = scalar.addf %total_g4, %zero : f32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%sum) : (f32) -> (f32) + scf.if %leader { + view.store %sum8, %sv[%p4, %local_row] : f32, view<5x[%row_tile]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %workitem = kernel.workitem.id : index + %publish_packets = index.div %row_tile, %c4 : index + %publish_token = index.div %workitem, %publish_packets : index + %publish_packet = index.rem %workitem, %publish_packets : index + %publish_add = index.mul %publish_packet, %c4 : index + %publish_channel = index.add %row_base, %publish_add : index + %valid_token = index.cmp ult, %publish_token, %tokens : index + %valid_channel = index.cmp ult, %publish_channel, %n : index + %publishes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %publishes { + %safe_token, %safe_count = index.assume %publish_token, %tokens [lt(%publish_token, %tokens)] : index, index + %stage_token = index.assume %safe_token [range(%safe_token, 0, 4)] : index + %row_limit = index.sub %row_tile, %c4 : index + %stage_add = index.assume %publish_add [range(%publish_add, 0, %row_limit), mul(%publish_add, 4)] : index + %stage_row_offset = index.mul %stage_token, %row_tile : index + %stage_origin0 = index.add %stage_row_offset, %stage_add : index + %stage_limit = index.sub %stage_elements, %c4 : index + %stage_origin = index.assume %stage_origin0 [range(%stage_origin0, 0, %stage_limit), mul(%stage_origin0, 4)] : index + %flat_stage = buffer.view %stage[%base] : buffer -> view<[%stage_elements]xf32> + %values = vector.load %flat_stage[%stage_origin] : view<[%stage_elements]xf32> -> vector<4xf32> + template.apply<@ggml.mul_mat_f32_f32_wmma.publish_vector4>(%publishes, %safe_token, %publish_channel, %safe_count, %n, %output_accumulation, %output_unary_op, %values, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (i1, index, index, index, index, index, index, vector<4xf32>, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + %rounding = index.sub %row_tile, %c1 : index + %rounded_n = index.add %n, %rounding : index + %output_tiles = index.div %rounded_n, %row_tile : index + template.apply<@ggml.mul_mat_f32_f32_wmma.finish_tile>(%tokens, %n, %output_tiles, %c1, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %c0, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +func.def inline @ggml_q4_prefetch_native_words(%weight: buffer, %row_group_block: index, %block: index, %row_lane: index, %payload_field: index, %payload_word: index) -> (vector<4xi32>, vector<2xi32>) { + %zero = index.constant 0 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %group_block = index.add %row_group_block, %block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %header_view = buffer.view %weight[%wb] : buffer -> view<4xi32> + %code_view = buffer.view %weight[%wc] : buffer -> view<8x64x4xi32> + %headers = vector.load %header_view[%zero] : view<4xi32> -> vector<4xi32> + %raw = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + func.return %headers, %raw : vector<4xi32>, vector<2xi32> +} + +// Native Q4_K bytes and Q8_1_x4 activations, with one 16-lane cohort per row. +// Retain the wrapper publication hooks so bias, residual and RMSNorm fusions +// share this contraction without a separate dispatch or intermediate output. +func.def pure inline @ggml_q4_q8_lowtoken_token_pairs(%capacity: index, %tokens: index, %k: index, %n: index) -> (i1) { + %c4 = index.constant 4 : index + %c4096 = index.constant 4096 : index + %capacity4 = index.cmp eq, %capacity, %c4 : index + %tokens4 = index.cmp eq, %tokens, %c4 : index + %four_tokens = scalar.andi %capacity4, %tokens4 : i1 + %long_input = index.cmp uge, %k, %c4096 : index + %wide_output = index.cmp uge, %n, %c4096 : index + %contracts = index.cmp ule, %n, %k : index + %wide = scalar.andi %long_input, %wide_output : i1 + %contracting = scalar.andi %wide, %contracts : i1 + %eligible = scalar.andi %four_tokens, %contracting : i1 + func.return %eligible : i1 +} + +func.def inline @ggml_q4_q8_lowtoken_adjacent_scale_min(%scale0: i32, %scale1: i32, %scale2: i32, %pair0: index) -> (i32, i32, i32, i32) { + %pair = index.assume %pair0 [range(%pair0, 0, 3)] : index + %two = index.constant 2 : index + %sixteen = index.constant 16 : index + %low = index.cmp ult, %pair, %two : index + %pair_lane = index.rem %pair, %two : index + %shift_index = index.mul %pair_lane, %sixteen : index + %shift = index.cast %shift_index : index to i32 + %two_i32 = scalar.constant 2 : i32 + %four_i32 = scalar.constant 4 : i32 + %eight_i32 = scalar.constant 8 : i32 + %low_mask = scalar.constant 3855 : i32 + %high_mask = scalar.constant 12336 : i32 + %byte_mask = scalar.constant 255 : i32 + %high_shift = scalar.addi %shift, %two_i32 : i32 + %min_shift = scalar.addi %shift, %four_i32 : i32 + %s_source = scf.select %low, %scale0, %scale2 : i32 + %m_source = scf.select %low, %scale1, %scale2 : i32 + %s_high_shift = scf.select %low, %shift, %high_shift : i32 + %m_low_shift = scf.select %low, %shift, %min_shift : i32 + %s_lo_raw = scalar.shrui %s_source, %shift : i32 + %s_lo = scalar.andi %s_lo_raw, %low_mask : i32 + %s_hi_raw = scalar.shrui %scale0, %s_high_shift : i32 + %s_hi = scalar.andi %s_hi_raw, %high_mask : i32 + %s = scalar.ori %s_lo, %s_hi : i32 + %m_lo_raw = scalar.shrui %m_source, %m_low_shift : i32 + %m_lo = scalar.andi %m_lo_raw, %low_mask : i32 + %m_hi_raw = scalar.shrui %scale1, %s_high_shift : i32 + %m_hi = scalar.andi %m_hi_raw, %high_mask : i32 + %m = scalar.ori %m_lo, %m_hi : i32 + %s0 = scalar.andi %s, %byte_mask : i32 + %m0 = scalar.andi %m, %byte_mask : i32 + %s1 = scalar.shrui %s, %eight_i32 : i32 + %m1 = scalar.shrui %m, %eight_i32 : i32 + func.return %s0, %m0, %s1, %m1 : i32, i32, i32, i32 +} + +func.def inline @ggml_q4_q8_lowtoken_dot8(%paired_tokens: i1, %weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) { + %value = scf.if %paired_tokens -> (f32) { + %zero = vector.constant 0 : vector<1xi32> + %wi0 = vector.extract %weight[0] : vector<2xi32> -> i32 + %ai0 = vector.extract %activation[0] : vector<2xi32> -> i32 + %ww0 = vector.splat %wi0 : vector<1xi32> + %aw0 = vector.splat %ai0 : vector<1xi32> + %w0 = vector.bitcast %ww0 : vector<1xi32> to vector<4xi8> + %a0 = vector.bitcast %aw0 : vector<1xi32> to vector<4xi8> + %d0 = vector.dot4i %w0, %a0, %zero : vector<4xi8>, vector<4xi8>, vector<1xi32> + %wi1 = vector.extract %weight[1] : vector<2xi32> -> i32 + %ai1 = vector.extract %activation[1] : vector<2xi32> -> i32 + %ww1 = vector.splat %wi1 : vector<1xi32> + %aw1 = vector.splat %ai1 : vector<1xi32> + %w1 = vector.bitcast %ww1 : vector<1xi32> to vector<4xi8> + %a1 = vector.bitcast %aw1 : vector<1xi32> to vector<4xi8> + %d1 = vector.dot4i %w1, %a1, %d0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %sum = vector.extract %d1[0] : vector<1xi32> -> i32 + %result = scalar.sitofp %sum : i32 to f32 + scf.yield %result : f32 + } else { + %original = func.call @ggml_dot_u8_s8_vector8_f32(%weight, %activation) : (vector<2xi32>, vector<2xi32>) -> (f32) + scf.yield %original : f32 + } + func.return %value : f32 +} + +func.def inline @ggml_mul_mat_q4_q8_1_x4_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 1, 5), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %paired_tokens = func.call pure @ggml_q4_q8_lowtoken_token_pairs(%capacity, %tokens, %k, %n) : (index, index, index, index) -> (i1) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %wg = index.add %channel_tile, %c0 : index + %c64_rows = index.constant 64 : index + %direct_rows = index.div %c8, %k_partitions : index + %row_tile = scf.select %staged, %c64_rows, %direct_rows : index + %row_base = index.mul %wg, %row_tile : index + %wave_rows = index.div %c4, %k_partitions : index + %wave_row = index.mul %wave, %wave_rows : index + %cohort_width = index.mul %c16, %k_partitions : index + %row_half = index.div %lane, %cohort_width : index + %cohort_lane = index.rem %lane, %cohort_width : index + %local_row = index.add %wave_row, %row_half : index + %row = index.add %row_base, %local_row : index + %valid_row = index.cmp ult, %row, %n : index + %cohort = index.div %lane, %c16 : index + %block_lane = index.rem %lane, %c16 : index + %pair = index.div %block_lane, %c4 : index + %packet = index.rem %block_lane, %c4 : index + %group0 = index.mul %pair, %c2 : index + %group1 = index.add %group0, %c1 : index + %word_page = index.mul %pair, %c8 : index + %word_add = index.mul %packet, %c2 : index + %word_index = index.add %word_page, %word_add : index + %meta_lane = index.cmp eq, %packet, %c0 : index + %blocks = index.div %k, %c256 : index + %partition_blocks = index.div %blocks, %k_partitions : index + %k_team = index.div %cohort_lane, %c16 : index + %k_begin = index.mul %k_team, %partition_blocks : index + %weight_row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %weight_row_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_groups = index.div %k, %c128 : index + %global_q8_row_bytes = index.scale %q8_groups, %block_bytes : index, offset -> offset + %slab_row_bytes = index.constant 2304 : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %q8_half = index.div %pair, %c2 : index + %q8_pair = index.rem %pair, %c2 : index + %q8_group0 = index.mul %q8_pair, %c2 : index + %q8_group1 = index.add %q8_group0, %c1 : index + %q8_words_base = index.mul %q8_group0, %c8 : index + %q8_word0 = index.add %q8_words_base, %word_add : index + %q8_word1 = index.add %q8_word0, %c8 : index + %q8_meta0 = index.mul %q8_group0, %c2 : index + %q8_meta1 = index.mul %q8_group1, %c2 : index + %q8_sum0 = index.add %q8_meta0, %c1 : index + %q8_sum1 = index.add %q8_meta1, %c1 : index + %c64 = index.constant 64 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %channel_tile, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %payload_field = index.div %word_index, %c4 : index + %payload_word = index.rem %word_index, %c4 : index + %input_na, %gate_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = index.constant 8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_header = vector.constant 0 : vector<4xi32> + %zero_raw = vector.constant 0 : vector<2xi32> + %prime_header, %prime_raw = scf.if %staged -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %c0, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %zero_header, %zero_raw : vector<4xi32>, vector<2xi32> + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_header, %last_raw = scf.for %block_base = [%c0 to %partition_blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_header = %prime_header : vector<4xi32>, %carry_raw = %prime_raw : vector<2xi32>) -> (f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32>) { + %block_in_slab = index.rem %block_base, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %new_slab = scalar.andi %staged, %slab_boundary : i1 + scf.if %new_slab { + func.call @ggml_q8_1_x4_stage_lowtoken_slab(%input_na, %q8_stage, %capacity, %block_base, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %block = index.add %block_base, %k_begin : index + %valid_block = index.cmp ult, %block, %blocks : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %group_block = index.add %row_group_block, %safe_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %g_header_view = buffer.view %gate_na[%wb] : buffer -> view<4xi32> + %g_code_view = buffer.view %gate_na[%wc] : buffer -> view<8x64x4xi32> + %g_headers = scf.if %staged -> (vector<4xi32>) { + scf.yield %carry_header : vector<4xi32> + } else { + %header = vector.load %g_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %header : vector<4xi32> + } + %g_dm_i32 = vector.extract %g_headers[0] : vector<4xi32> -> i32 + %g_dm_word = vector.splat %g_dm_i32 : vector<1xi32> + %g_dm = vector.bitcast %g_dm_word : vector<1xi32> to vector<2xf16> + %g_d0h = vector.extract %g_dm[0] : vector<2xf16> -> f16 + %g_d1h = vector.extract %g_dm[1] : vector<2xf16> -> f16 + %g_d0 = scalar.extf %g_d0h : f16 to f32 + %g_d1 = scalar.extf %g_d1h : f16 to f32 + %g_s0 = vector.extract %g_headers[1] : vector<4xi32> -> i32 + %g_s1 = vector.extract %g_headers[2] : vector<4xi32> -> i32 + %g_s2 = vector.extract %g_headers[3] : vector<4xi32> -> i32 + %pair_scale0, %pair_min0, %pair_scale1, %pair_min1 = scf.if %paired_tokens -> (i32, i32, i32, i32) { + %s0, %m0, %s1, %m1 = func.call @ggml_q4_q8_lowtoken_adjacent_scale_min(%g_s0, %g_s1, %g_s2, %pair) : (i32, i32, i32, index) -> (i32, i32, i32, i32) + scf.yield %s0, %m0, %s1, %m1 : i32, i32, i32, i32 + } else { + %zero_i32 = scalar.constant 0 : i32 + scf.yield %zero_i32, %zero_i32, %zero_i32, %zero_i32 : i32, i32, i32, i32 + } + %g_scale0i, %g_min0i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale0, %pair_min0 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale0f = scalar.uitofp %g_scale0i : i32 to f32 + %g_min0f = scalar.uitofp %g_min0i : i32 to f32 + %g_ws0 = scalar.mulf %g_d0, %g_scale0f : f32 + %g_wm0 = scalar.mulf %g_d1, %g_min0f : f32 + %g_scale1i, %g_min1i = scf.if %paired_tokens -> (i32, i32) { + scf.yield %pair_scale1, %pair_min1 : i32, i32 + } else { + %scale, %minimum = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + scf.yield %scale, %minimum : i32, i32 + } + %g_scale1f = scalar.uitofp %g_scale1i : i32 to f32 + %g_min1f = scalar.uitofp %g_min1i : i32 to f32 + %g_ws1 = scalar.mulf %g_d0, %g_scale1f : f32 + %g_wm1 = scalar.mulf %g_d1, %g_min1f : f32 + %g_raw = scf.if %staged -> (vector<2xi32>) { + scf.yield %carry_raw : vector<2xi32> + } else { + %raw = vector.load %g_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %raw : vector<2xi32> + } + %g_mask = vector.constant 252645135 : vector<2xi32> + %g_shift = vector.constant 4 : vector<2xi32> + %g_lo = vector.andi %g_raw, %g_mask : vector<2xi32> + %g_high = vector.shrui %g_raw, %g_shift : vector<2xi32> + %g_hi = vector.andi %g_high, %g_mask : vector<2xi32> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_header, %next_raw = scf.if %has_next -> (vector<4xi32>, vector<2xi32>) { + %header, %raw = func.call @ggml_q4_prefetch_native_words(%gate_na, %row_group_block, %next_block, %row_lane, %payload_field, %payload_word) : (buffer, index, index, index, index, index) -> (vector<4xi32>, vector<2xi32>) + scf.yield %header, %raw : vector<4xi32>, vector<2xi32> + } else { + scf.yield %carry_header, %carry_raw : vector<4xi32>, vector<2xi32> + } + %input_block = scf.select %staged, %block_in_slab, %safe_block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %q8_half : index + %q8_block_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next_g0 = scf.if %valid0 -> (f32) { + %tb = index.scale %t0, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ad0h = view.load %ds_view[%q8_meta0] : view<8xf16> -> f16 + %ad1h = view.load %ds_view[%q8_meta1] : view<8xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = view.load %ds_view[%q8_sum0] : view<8xf16> -> f16 + %as1h = view.load %ds_view[%q8_sum1] : view<8xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g0, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next_g1 = scf.if %valid1 -> (f32) { + %tb = index.scale %t1, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g1, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + scf.if %paired_tokens { + scf.schedule.fence + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next_g2 = scf.if %valid2 -> (f32) { + %tb = index.scale %t2, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g2, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next_g3 = scf.if %valid3 -> (f32) { + %tb = index.scale %t3, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g3, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next_g4 = scf.if %valid4 -> (f32) { + %tb = index.scale %t4, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %a0 = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %q_view[%q8_word1] : view<32xi32> -> vector<2xi32> + %ds = vector.load %ds_view[%q8_meta0] : view<8xf16> -> vector<4xf16> + %ad0h = vector.extract %ds[0] : vector<4xf16> -> f16 + %ad1h = vector.extract %ds[2] : vector<4xf16> -> f16 + %ad0 = scalar.extf %ad0h : f16 to f32 + %ad1 = scalar.extf %ad1h : f16 to f32 + %as0h = vector.extract %ds[1] : vector<4xf16> -> f16 + %as1h = vector.extract %ds[3] : vector<4xf16> -> f16 + %as0_raw = scalar.extf %as0h : f16 to f32 + %as1_raw = scalar.extf %as1h : f16 to f32 + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_lo, %a0) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_q4_q8_lowtoken_dot8(%paired_tokens, %g_hi, %a1) : (i1, vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g4, %g_valid : f32 + scf.yield %g_updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_header, %next_raw : f32, f32, f32, f32, f32, vector<4xi32>, vector<2xi32> + } + %combined0 = func.call @ggml_lowtoken_combine_k_partitions(%total_g0, %k_partitions) : (f32, index) -> (f32) + %combined1 = func.call @ggml_lowtoken_combine_k_partitions(%total_g1, %k_partitions) : (f32, index) -> (f32) + %combined2 = func.call @ggml_lowtoken_combine_k_partitions(%total_g2, %k_partitions) : (f32, index) -> (f32) + %combined3 = func.call @ggml_lowtoken_combine_k_partitions(%total_g3, %k_partitions) : (f32, index) -> (f32) + %combined4 = func.call @ggml_lowtoken_combine_k_partitions(%total_g4, %k_partitions) : (f32, index) -> (f32) + func.call @ggml_mul_mat_q8_lowtoken_publish(%tokens, %n, %row_tile, %row_base, %local_row, %cohort_lane, %channel_tile, %output_accumulation, %output_unary_op, %epsilon, %combined0, %combined1, %combined2, %combined3, %combined4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, index, index, f32, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +// Native Q6_K carrier with signed I8 dot products and F32 scaling. +func.def inline @ggml_q6_prefetch_native_words(%weight_na: buffer, %row_group_block: index, %block: index, %row_lane: index, %half16: index, %group: index, %half128: index, %ql_side: index, %packet: index) -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %c2 = index.constant 2 : index + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %group_block = index.add %row_group_block, %block : index + %block_base = index.scale %group_block, %record_bytes : index, offset -> offset + %scale_base = index.add %block_base, %scale_add : offset + %raw_base = index.add %block_base, %raw_add : offset + %d_view = buffer.view %weight_na[%block_base] : buffer -> view<64xf16> + %s_view = buffer.view %weight_na[%scale_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_na[%raw_base] : buffer -> view<2x3x64x8xi32> + %dh = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %si = view.load %s_view[%row_lane, %half16, %group] : view<64x2x8xi8> -> i8 + %ql = vector.load %raw_view[%half128, %ql_side, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + %qh = vector.load %raw_view[%half128, %c2, %row_lane, %packet] : view<2x3x64x8xi32> -> vector<4xi32> + func.return %dh, %si, %ql, %qh : f16, i8, vector<4xi32>, vector<4xi32> +} + +func.def inline @ggml_q6_q8_lowtoken_activation_values(%pipelined: i1, %token: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index) -> (vector<16xi8>, f32) { + %code_offset = index.constant 16 : offset + %a, %ad = scf.if %pipelined -> (vector<16xi8>, f32) { + %base = index.constant 0 : offset + %c4 = index.constant 4 : index + %q8_block = index.assume %q8_block0 [range(%q8_block0, 0, 7)] : index + %word = index.add %q8_word0, %c4 : index + %q_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x36xi32> + %ds_view = buffer.view %q8_input[%base] : buffer -> view<2x[%capacity]x8x72xf16> + %words = vector.load %q_view[%current_slot, %token, %q8_block, %word] : view<2x[%capacity]x8x36xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%current_slot, %token, %q8_block, %q8_meta] : view<2x[%capacity]x8x72xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } else { + %tb = index.scale %token, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %qc = index.add %qb, %code_offset : offset + %q_view = buffer.view %q8_input[%qc] : buffer -> view<32xi32> + %ds_view = buffer.view %q8_input[%qb] : buffer -> view<8xf16> + %words = vector.load %q_view[%q8_word0] : view<32xi32> -> vector<4xi32> + %a = vector.bitcast %words : vector<4xi32> to vector<16xi8> + %scale = view.load %ds_view[%q8_meta] : view<8xf16> -> f16 + %ad = scalar.extf %scale : f16 to f32 + scf.yield %a, %ad : vector<16xi8>, f32 + } + func.return %a, %ad : vector<16xi8>, f32 +} + +func.def inline @ggml_q6_q8_lowtoken_dot16(%pipelined: i1, %w: vector<16xi8>, %a: vector<16xi8>) -> (i32) { + %sum = scf.if %pipelined -> (i32) { + %di0 = vector.constant 0 : vector<1xi32> + %dw0 = vector.slice %w[0] : vector<16xi8> -> vector<4xi8> + %da0 = vector.slice %a[0] : vector<16xi8> -> vector<4xi8> + %di1 = vector.dot4i %dw0, %da0, %di0 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw1 = vector.slice %w[4] : vector<16xi8> -> vector<4xi8> + %da1 = vector.slice %a[4] : vector<16xi8> -> vector<4xi8> + %di2 = vector.dot4i %dw1, %da1, %di1 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw2 = vector.slice %w[8] : vector<16xi8> -> vector<4xi8> + %da2 = vector.slice %a[8] : vector<16xi8> -> vector<4xi8> + %di3 = vector.dot4i %dw2, %da2, %di2 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %dw3 = vector.slice %w[12] : vector<16xi8> -> vector<4xi8> + %da3 = vector.slice %a[12] : vector<16xi8> -> vector<4xi8> + %di4 = vector.dot4i %dw3, %da3, %di3 : vector<4xi8>, vector<4xi8>, vector<1xi32> + %chain_sum = vector.extract %di4[0] : vector<1xi32> -> i32 + scf.yield %chain_sum : i32 + } else { + %init = vector.constant 0 : vector<4xi32> + %dot = vector.dot4i %w, %a, %init : vector<16xi8>, vector<16xi8>, vector<4xi32> + %seed = scalar.constant 0 : i32 + %direct_sum = vector.reduce %dot, %seed : vector<4xi32>, i32 + scf.yield %direct_sum : i32 + } + func.return %sum : i32 +} + +func.def inline @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input: buffer, %capacity0: index, %slot_begin: index, %blocks: index, %global_row_bytes: offset) -> (vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c18 = index.constant 18 : index + %c72 = index.constant 72 : index + %packet_bytes = index.constant 16 : offset + %block_bytes = index.constant 288 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %copy_packet = index.rem %thread, %total_packets : index + %copy_token = index.div %copy_packet, %c72 : index + %copy_inner = index.rem %copy_packet, %c72 : index + %remaining = index.sub %blocks, %slot_begin : index + %full = index.cmp uge, %remaining, %c4 : index + %active0 = scf.select %full, %c4, %remaining : index + %active = index.assume %active0 [range(%active0, 1, 4)] : index + %active_packets = index.mul %active, %c18 : index + %copy_valid = index.cmp ult, %copy_inner, %active_packets : index + %safe_inner = scf.select %copy_valid, %copy_inner, %c0 : index + %copy_row = index.scale %copy_token, %global_row_bytes : index, offset -> offset + %slot_bytes = index.scale %slot_begin, %block_bytes : index, offset -> offset + %copy_bytes = index.scale %safe_inner, %packet_bytes : index, offset -> offset + %copy_base0 = index.add %copy_row, %slot_bytes : offset + %copy_base = index.add %copy_base0, %copy_bytes : offset + %source = buffer.view %input[%copy_base] : buffer -> view<4xi32> + %values = vector.load %source[%c0] : view<4xi32> -> vector<4xi32> + func.return %values : vector<4xi32> +} +func.def inline @ggml_q8_1_x4_publish_lowtoken_half_slab(%stage: buffer, %capacity0: index, %slot0: index, %values: vector<4xi32>) { + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %slot = index.assume %slot0 [range(%slot0, 0, 1)] : index + %c4 = index.constant 4 : index + %c72 = index.constant 72 : index + %base = index.constant 0 : offset + %thread = kernel.workitem.id : index + %total_packets = index.mul %capacity, %c72 : index + %thread_valid = index.cmp ult, %thread, %total_packets : index + scf.if %thread_valid { + %copy_token0 = index.div %thread, %c72 : index + %copy_token = index.assume %copy_token0 [range(%copy_token0, 0, 4), lt(%copy_token0, %capacity)] : index + %copy_inner = index.rem %thread, %c72 : index + %origin = index.mul %copy_inner, %c4 : index + %destination = buffer.view %stage[%base] : buffer -> view<2x[%capacity]x288xi32> + vector.store %values, %destination[%slot, %copy_token, %origin] : vector<4xi32>, view<2x[%capacity]x288xi32> + } + func.return +} + + +func.def inline @ggml_q6_q8_lowtoken_accumulate(%pipelined: i1, %tokens: index, %q8_row_bytes: offset, %q8_block_add: offset, %capacity: index, %current_slot: index, %q8_block0: index, %q8_input: buffer, %q8_word0: index, %q8_meta: index, %w: vector<16xi8>, %ws: f32, %acc_g0: f32, %acc_g1: f32, %acc_g2: f32, %acc_g3: f32, %acc_g4: f32) -> (f32, f32, f32, f32, f32) { + %code_offset = index.constant 16 : offset + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %computed_g0 = scf.if %valid0 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t0, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g0 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g0 : f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %computed_g1 = scf.if %valid1 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t1, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g1 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g1 : f32 + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %computed_g2 = scf.if %valid2 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t2, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g2 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g2 : f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %computed_g3 = scf.if %valid3 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t3, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g3 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g3 : f32 + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %computed_g4 = scf.if %valid4 -> (f32) { + %a, %ad = func.call @ggml_q6_q8_lowtoken_activation_values(%pipelined, %t4, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block0, %q8_input, %q8_word0, %q8_meta) : (i1, index, offset, offset, index, index, index, buffer, index, index) -> (vector<16xi8>, f32) + %sum = func.call @ggml_q6_q8_lowtoken_dot16(%pipelined, %w, %a) : (i1, vector<16xi8>, vector<16xi8>) -> (i32) + %sumf = scalar.sitofp %sum : i32 to f32 + %scale = scalar.mulf %ws, %ad : f32 + %updated = scalar.fmaf %sumf, %scale, %acc_g4 : f32 + scf.yield %updated : f32 + } else { + scf.yield %acc_g4 : f32 + } + func.return %computed_g0, %computed_g1, %computed_g2, %computed_g3, %computed_g4 : f32, f32, f32, f32, f32 +} + +func.def inline @ggml_mul_mat_q6_q8_1_x4_lowtoken_dot(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %capacity = config.get @ggml.workload.token_capacity : index + %staged_rows, %k_partitions = func.call pure @ggml_q8_lowtoken_geometry(%capacity, %input_size0, %output_size0, %weight_format) : (index, index, index, index) -> (index, index) + %no_staged_rows = index.constant 0 : index + %staged = index.cmp ugt, %staged_rows, %no_staged_rows : index + %tokens = index.assume %token_count [range(%token_count, 1, 5), le(%token_count, %capacity)] : index + %k = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %n = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %row_tile = scf.select %staged, %c32, %c8 : index + %row_base = index.mul %channel_tile, %row_tile : index + %wave_row = index.mul %wave, %c4 : index + %cohort = index.div %lane, %c16 : index + %local_row = index.add %wave_row, %cohort : index + %row = index.add %row_base, %local_row : index + %block_lane = index.rem %lane, %c16 : index + %group = index.div %block_lane, %c2 : index + %half16 = index.rem %block_lane, %c2 : index + %half128 = index.div %group, %c4 : index + %group_in_half = index.rem %group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift0 = index.mul %nibble, %c4 : index + %nibble_shift1 = index.cast %nibble_shift0 : index to i32 + %nibble_shift = vector.splat %nibble_shift1 : vector<4xi32> + %high_shift0 = index.mul %group_in_half, %c2 : index + %high_shift1 = index.cast %high_shift0 : index to i32 + %high_shift = vector.splat %high_shift1 : vector<4xi32> + %packet = index.mul %half16, %c4 : index + %q8_group_words = index.mul %group_in_half, %c8 : index + %q8_word0 = index.add %q8_group_words, %packet : index + %q8_meta = index.mul %group_in_half, %c2 : index + %blocks = index.div %k, %c256 : index + %q8_groups = index.div %k, %c128 : index + %q8_block_bytes = index.constant 144 : offset + %global_q8_row_bytes = index.scale %q8_groups, %q8_block_bytes : index, offset -> offset + %native_block_bytes = index.constant 210 : index + %cache_bytes = index.constant 33554432 : index + %weight_blocks = index.mul %blocks, %n : index + %weight_bytes = index.mul %weight_blocks, %native_block_bytes : index + %exceeds_cache = index.cmp uge, %weight_bytes, %cache_bytes : index + %pipelined = scalar.andi %staged, %exceeds_cache : i1 + %full_slab_row_bytes = index.constant 2304 : offset + %half_slab_row_bytes = index.constant 1152 : offset + %slab_row_bytes = scf.select %pipelined, %half_slab_row_bytes, %full_slab_row_bytes : offset + %q8_row_bytes = scf.select %staged, %slab_row_bytes, %global_q8_row_bytes : offset + %code_offset = index.constant 16 : offset + %record_bytes = index.constant 13440 : offset + %scale_add = index.constant 128 : offset + %raw_add = index.constant 1152 : offset + %staged_row_group = index.div %channel_tile, %c2 : index + %direct_row_group = index.div %row, %c64 : index + %row_group = scf.select %staged, %staged_row_group, %direct_row_group : index + %row_lane = index.rem %row, %c64 : index + %row_group_block = index.mul %row_group, %blocks : index + %input_na, %weight_na = buffer.assume.noalias %input, %weight : buffer, buffer + %slab_blocks = scf.select %pipelined, %c4, %c8 : index + %full_stage_bytes = index.constant 11520 : offset + %q8_stage_bytes = scf.select %staged, %full_stage_bytes, %base : offset + %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer + %q8_input = scf.select %staged, %q8_stage, %input_na : buffer + %zero_half = scalar.constant 0.0 : f16 + %zero_code = scalar.constant 0 : i8 + %zero_words = vector.constant 0 : vector<4xi32> + %prime_d, %prime_s, %prime_ql, %prime_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %pd, %ps, %pql, %pqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %c0, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %pd, %ps, %pql, %pqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %zero_half, %zero_code, %zero_words, %zero_words : f16, i8, vector<4xi32>, vector<4xi32> + } + scf.if %pipelined { + %prime_q8 = func.call @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input_na, %capacity, %c0, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + func.call @ggml_q8_1_x4_publish_lowtoken_half_slab(%q8_stage, %capacity, %c0, %prime_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %last_d, %last_s, %last_ql, %last_qh = scf.for %block = [%c0 to %blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %carry_d = %prime_d : f16, %carry_s = %prime_s : i8, %carry_ql = %prime_ql : vector<4xi32>, %carry_qh = %prime_qh : vector<4xi32>) -> (f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32>) { + %block_in_slab = index.rem %block, %slab_blocks : index + %slab_boundary = index.cmp eq, %block_in_slab, %c0 : index + %direct_stage = scalar.andi %staged, %slab_boundary : i1 + %true = scalar.constant true : i1 + %not_pipelined = scalar.xori %pipelined, %true : i1 + %stage_current = scalar.andi %direct_stage, %not_pipelined : i1 + scf.if %stage_current { + func.call @ggml_q8_1_x4_stage_lowtoken_slab(%input_na, %q8_stage, %capacity, %block, %blocks, %global_q8_row_bytes) : (buffer, buffer, index, index, index, offset) + } + %slab_number = index.div %block, %slab_blocks : index + %current_slot = index.rem %slab_number, %c2 : index + %following_slab = index.add %slab_number, %c1 : index + %following_begin = index.mul %following_slab, %slab_blocks : index + %following_exists = index.cmp ult, %following_begin, %blocks : index + %next_slot = index.rem %following_slab, %c2 : index + %staged_future = scalar.andi %pipelined, %following_exists : i1 + %current_d, %current_s, %current_ql, %current_qh = scf.if %staged -> (f16, i8, vector<4xi32>, vector<4xi32>) { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + %wd, %ws8, %wql, %wqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } + %d = scalar.extf %current_d : f16 to f32 + %sf = scalar.sitofp %current_s : i8 to f32 + %ws = scalar.mulf %d, %sf : f32 + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %four = vector.constant 4 : vector<4xi32> + %ql_shifted = vector.shrui %current_ql, %nibble_shift : vector<4xi32> + %low = vector.andi %ql_shifted, %nibble_mask : vector<4xi32> + %qh_shifted = vector.shrui %current_qh, %high_shift : vector<4xi32> + %high0 = vector.andi %qh_shifted, %high_mask : vector<4xi32> + %high = vector.shli %high0, %four : vector<4xi32> + %code = vector.ori %low, %high : vector<4xi32> + // Guard bits prevent subtraction from borrowing across packed bytes. + %guard_mask = vector.constant -2139062144 : vector<4xi32> + %center = vector.constant 538976288 : vector<4xi32> + %guarded = vector.ori %code, %guard_mask : vector<4xi32> + %centered = vector.subi %guarded, %center : vector<4xi32> + %signed_code = vector.xori %centered, %guard_mask : vector<4xi32> + %w = vector.bitcast %signed_code : vector<4xi32> to vector<16xi8> + %next_block = index.add %block, %c1 : index + %next_exists = index.cmp ult, %next_block, %blocks : index + %has_next = scalar.andi %staged, %next_exists : i1 + %next_d, %next_s, %next_ql, %next_qh = scf.if %has_next -> (f16, i8, vector<4xi32>, vector<4xi32>) { + %wd, %ws8, %wql, %wqh = func.call @ggml_q6_prefetch_native_words(%weight_na, %row_group_block, %next_block, %row_lane, %half16, %group, %half128, %ql_side, %packet) : (buffer, index, index, index, index, index, index, index, index) -> (f16, i8, vector<4xi32>, vector<4xi32>) + scf.yield %wd, %ws8, %wql, %wqh : f16, i8, vector<4xi32>, vector<4xi32> + } else { + scf.yield %carry_d, %carry_s, %carry_ql, %carry_qh : f16, i8, vector<4xi32>, vector<4xi32> + } + %input_block = scf.select %staged, %block_in_slab, %block : index + %q8_block0 = index.mul %input_block, %c2 : index + %q8_block = index.add %q8_block0, %half128 : index + %q8_local_block_add = index.scale %q8_block, %q8_block_bytes : index, offset -> offset + %slot_row = index.mul %current_slot, %capacity : index + %slot_byte_add0 = index.scale %slot_row, %slab_row_bytes : index, offset -> offset + %slot_byte_add = scf.select %pipelined, %slot_byte_add0, %base : offset + %q8_block_add = index.add %slot_byte_add, %q8_local_block_add : offset + %slot_complete = index.cmp eq, %block_in_slab, %c3 : index + %publish_next = scalar.andi %staged_future, %slot_complete : i1 + %next_g0, %next_g1, %next_g2, %next_g3, %next_g4 = scf.if %publish_next -> (f32, f32, f32, f32, f32) { + %next_q8 = func.call @ggml_q8_1_x4_prefetch_lowtoken_half_slab(%input_na, %capacity, %following_begin, %blocks, %global_q8_row_bytes) : (buffer, index, index, index, offset) -> (vector<4xi32>) + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_q6_q8_lowtoken_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + func.call @ggml_q8_1_x4_publish_lowtoken_half_slab(%q8_stage, %capacity, %next_slot, %next_q8) : (buffer, index, index, vector<4xi32>) + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } else { + %v0, %v1, %v2, %v3, %v4 = func.call @ggml_q6_q8_lowtoken_accumulate(%pipelined, %tokens, %q8_row_bytes, %q8_block_add, %capacity, %current_slot, %q8_block, %q8_input, %q8_word0, %q8_meta, %w, %ws, %acc_g0, %acc_g1, %acc_g2, %acc_g3, %acc_g4) : (i1, index, offset, offset, index, index, index, buffer, index, index, vector<16xi8>, f32, f32, f32, f32, f32, f32) -> (f32, f32, f32, f32, f32) + scf.yield %v0, %v1, %v2, %v3, %v4 : f32, f32, f32, f32, f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_d, %next_s, %next_ql, %next_qh : f32, f32, f32, f32, f32, f16, i8, vector<4xi32>, vector<4xi32> + } + func.call @ggml_mul_mat_q8_lowtoken_publish(%tokens, %n, %row_tile, %row_base, %local_row, %block_lane, %channel_tile, %output_accumulation, %output_unary_op, %epsilon, %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, index, index, f32, f32, f32, f32, f32, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + func.return +} + +template.def<@ggml.mul_mat_f32_f32_wmma.core> device @ggml_mul_mat_f32_f32_wmma_core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: 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) { + %token_capacity = config.get @ggml.workload.token_capacity : index + %lowtoken = func.call pure @ggml_mul_mat_uses_lowtoken_dot(%token_capacity, %input_size0) : (index, index) -> (i1) + %activation_format = config.get @ggml.mul_mat.activation_format : index + %q8_format = index.constant 9 : index + %uses_q8 = index.cmp eq, %activation_format, %q8_format : index + scf.if %lowtoken { + scf.if %uses_q8 { + %q6_format = index.constant 46 : index + %uses_q6 = index.cmp eq, %weight_format, %q6_format : index + scf.if %uses_q6 { + func.call @ggml_mul_mat_q6_q8_1_x4_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } else { + func.call @ggml_mul_mat_q4_q8_1_x4_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } else { + func.call @ggml_mul_mat_f32_f32_lowtoken_dot(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + } else { + func.call @ggml_mul_mat_f32_f32_wmma_tiled(%weight_format, %token_count, %input_size0, %output_size0, %output_accumulation, %output_unary_op, %epsilon, %channel_tile, %token_tile, %input, %weight, %bias, %output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom new file mode 100644 index 000000000000..07f2ef3fe8aa --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom @@ -0,0 +1,780 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.binary_f32.apply_vector8(%op: index, %lhs: vector<8xf32>, %rhs: vector<8xf32>) -> (vector<8xf32>) + +func.decl @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) +func.decl @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) + +func.def inline @ggml_f16_prefill_load_pair(%input: buffer, %token_count: index, %input_size: index, %token_origin: index, %k_origin: index) -> (vector<8xf16>, vector<8xf16>) { + %c0_offset = index.constant 0 : offset + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c1024 = index.constant 1024 : index + %workgroup_size = kernel.workgroup.size : index + %larger_tile = index.cmp eq, %workgroup_size, %c1024 : index + %row_packets = scf.select %larger_tile, %c4, %c2 : index + %packet_width = scf.select %larger_tile, %c8, %c16 : index + %workitem = kernel.workitem.id : index + %row0 = index.div %workitem, %row_packets : index + %row = index.assume %row0 [range(%row0, 0, 255)] : index + %packet = index.rem %workitem, %row_packets : index + %packet_k = index.mul %packet, %packet_width : index + %token = index.add %token_origin, %row : index + %input_k0 = index.add %k_origin, %packet_k : index + %input_k1 = index.add %input_k0, %c8 : index + %view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf16> + %valid = index.cmp ult, %token, %token_count : index + %zero = vector.constant 0.0 : vector<8xf16> + %lo, %hi = scf.if %valid -> (vector<8xf16>, vector<8xf16>) { + %bounded_token, %bounded_count = index.assume %token, %token_count [lt(%token, %token_count)] : index, index + %v0 = vector.load %view[%bounded_token, %input_k0] : view<[%token_count]x[%input_size]xf16> -> vector<8xf16> + %v1 = scf.if %larger_tile -> (vector<8xf16>) { + scf.yield %zero : vector<8xf16> + } else { + %loaded = vector.load %view[%bounded_token, %input_k1] : view<[%token_count]x[%input_size]xf16> -> vector<8xf16> + scf.yield %loaded : vector<8xf16> + } + scf.yield %v0, %v1 : vector<8xf16>, vector<8xf16> + } else { + scf.yield %zero, %zero : vector<8xf16>, vector<8xf16> + } + func.return %lo, %hi : vector<8xf16>, vector<8xf16> +} + + +func.def inline @ggml_f16_prefill_stage_mma(%paired: i1, %packed_input: i1, %bounded_token_count: index, %bounded_input_size: index, %token_tile_origin: index, %quant_block: index, %weight_group_base: index, %quant_group: index, %input_noalias: buffer, %weight_stage: buffer, %activation_stage: buffer, %stage_offset: offset, %pipeline_weights: i1, %acc00: vector<8xf32>, %acc01: vector<8xf32>, %acc02: vector<8xf32>, %acc03: vector<8xf32>, %acc10: vector<8xf32>, %acc11: vector<8xf32>, %acc12: vector<8xf32>, %acc13: vector<8xf32>, %current_a0: vector<8xf16>, %current_a1: vector<8xf16>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) { + %layout_true = scalar.constant true : i1 + %sync_weights = scalar.xori %pipeline_weights, %layout_true : i1 + %staged_input = scalar.xori %packed_input, %layout_true : i1 + %workgroup_size = kernel.workgroup.size : index + %large_workgroup = index.constant 512 : index + %wide_tile = index.cmp uge, %workgroup_size, %large_workgroup : index + %can_prefetch_a = index.cmp eq, %workgroup_size, %large_workgroup : index + %prefetch_a = scalar.andi %can_prefetch_a, %staged_input : i1 + %largest_workgroup = index.constant 1024 : index + %wg1024 = index.cmp eq, %workgroup_size, %largest_workgroup : index + %larger_tile = scalar.andi %wg1024, %staged_input : i1 + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c136 = index.constant 136 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c512 = index.constant 512 : index + %c0_offset = index.constant 0 : offset + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %true = scalar.constant true : i1 + %not_wide_tile = scalar.xori %wide_tile, %true : i1 + %zero_f16 = vector.constant 0.0 : vector<16xf16> + %weight_stride = scf.select %wide_tile, %c72, %c136 : index + %ordinary_weight_rows = scf.select %wide_tile, %c128, %c64 : index + %staged_weight_rows = scf.select %larger_tile, %c256, %ordinary_weight_rows : index + %packed_weight_rows = scf.select %wg1024, %c128, %c64 : index + %weight_rows = scf.select %packed_input, %packed_weight_rows, %staged_weight_rows : index + %weight_window = scf.select %wide_tile, %c64, %c128 : index + %input_f16_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<256x40xf16> + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %ordinary_subgroup_limit = scf.select %wide_tile, %c15, %c7 : index + %subgroup_limit = scf.select %wg1024, %c31, %ordinary_subgroup_limit : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, %subgroup_limit)] : index + %token_subgroups = scf.select %packed_input, %c16, %c8 : index + %token_subgroup = index.rem %subgroup, %token_subgroups : index + %subgroup_token_add = index.mul %token_subgroup, %c32 : index + %channel_subgroup = index.div %subgroup, %token_subgroups : index + %subgroup_channel_stride = scf.select %paired, %c32, %c64 : index + %subgroup_channel_add = index.mul %channel_subgroup, %subgroup_channel_stride : index + %half_weight_rows = index.div %weight_rows, %c2 : index + %peer_channel_delta = scf.select %paired, %half_weight_rows, %c32 : index + %rhs_column0 = index.add %subgroup_channel_add, %c0 : index + %rhs_column1 = index.add %rhs_column0, %c16 : index + %rhs_column2 = index.add %rhs_column0, %peer_channel_delta : index + %rhs_column3 = index.add %rhs_column2, %c16 : index + %activation_row_packets = scf.select %larger_tile, %c4, %c2 : index + %activation_packet_width = scf.select %larger_tile, %c8, %c16 : index + %activation_load_row0 = index.div %workitem, %activation_row_packets : index + %activation_row_limit = scf.select %wide_tile, %c255, %c127 : index + %activation_load_row = index.assume %activation_load_row0 [range(%activation_load_row0, 0, %activation_row_limit)] : index + %activation_packet = index.rem %workitem, %activation_row_packets : index + %activation_load_k = index.mul %activation_packet, %activation_packet_width : index + %second_token_tile_base = index.add %token_tile_origin, %c128 : index + %synchronous_a = scalar.xori %prefetch_a, %true : i1 + %k_origin = scf.if %synchronous_a -> (index) { + %block_k_base = index.mul %quant_block, %c256 : index + %absolute_group = index.add %weight_group_base, %quant_group : index + %group_k_add = index.mul %absolute_group, %c32 : index + %ordinary_k = index.add %block_k_base, %group_k_add : index + scf.yield %ordinary_k : index + } else { + scf.yield %c0 : index + } + scf.if %staged_input { + %token0 = index.add %token_tile_origin, %activation_load_row : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + scf.if %larger_tile { + %loaded, %unused = func.call @ggml_f16_prefill_load_pair(%input_noalias, %bounded_token_count, %bounded_input_size, %token_tile_origin, %k_origin) : (buffer, index, index, index, index) -> (vector<8xf16>, vector<8xf16>) + vector.store %loaded, %activation_stage_physical_view[%activation_load_row, %activation_load_k] : vector<8xf16>, view<256x40xf16> + scf.yield + } else { + %activation_values0 = scf.if %prefetch_a -> (vector<16xf16>) { + %values = vector.concat<0> %current_a0, %current_a1 : vector<8xf16>, vector<8xf16> -> vector<16xf16> + scf.yield %values : vector<16xf16> + } else { + %loaded_values = scf.if %valid_token0 -> (vector<16xf16>) { + %bounded_token, %input_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_k = index.add %k_origin, %activation_load_k : index + %loaded = vector.load %input_f16_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf16> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } else { + scf.yield %zero_f16 : vector<16xf16> + } + scf.yield %loaded_values : vector<16xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_load_row, %activation_load_k] : vector<16xf16>, view<256x40xf16> + scf.yield + } + scf.if %not_wide_tile { + %token1 = index.add %second_token_tile_base, %activation_load_row : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %activation_values1 = scf.if %valid_token1 -> (vector<16xf16>) { + %bounded_token, %input_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index + %input_k = index.add %k_origin, %activation_load_k : index + %loaded = vector.load %input_f16_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf16> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } else { + scf.yield %zero_f16 : vector<16xf16> + } + %activation_second_row = index.add %activation_load_row, %c128 : index + vector.store %activation_values1, %activation_stage_physical_view[%activation_second_row, %activation_load_k] : vector<16xf16>, view<256x40xf16> + scf.yield + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } else { + %first_group = index.cmp eq, %quant_group, %c0 : index + %publish_stage = scalar.andi %first_group, %sync_weights : i1 + scf.if %publish_stage { + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + %next_a0, %next_a1 = scf.if %prefetch_a -> (vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %absolute_group = index.add %weight_group_base, %quant_group : index + %group_k_add = index.mul %absolute_group, %c32 : index + %wide_k = index.add %block_k_base, %group_k_add : index + %next_k = index.add %wide_k, %c32 : index + %has_next = index.cmp ult, %next_k, %bounded_input_size : index + %next_input_k = scf.select %has_next, %next_k, %wide_k : index + %a0, %a1 = func.call @ggml_f16_prefill_load_pair(%input_noalias, %bounded_token_count, %bounded_input_size, %token_tile_origin, %next_input_k) : (buffer, index, index, index, index) -> (vector<8xf16>, vector<8xf16>) + scf.yield %a0, %a1 : vector<8xf16>, vector<8xf16> + } else { + %zero = vector.constant 0.0 : vector<8xf16> + scf.yield %zero, %zero : vector<8xf16>, vector<8xf16> + } + %serial_rhs = scalar.ori %larger_tile, %packed_input : i1 + %next00, %next01, %next02, %next03, %next10, %next11, %next12, %next13 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf32>, %half_acc01 = %acc01 : vector<8xf32>, %half_acc02 = %acc02 : vector<8xf32>, %half_acc03 = %acc03 : vector<8xf32>, %half_acc10 = %acc10 : vector<8xf32>, %half_acc11 = %acc11 : vector<8xf32>, %half_acc12 = %acc12 : vector<8xf32>, %half_acc13 = %acc13 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) unroll { + %lhs_m0 = index.add %subgroup_token_add, %c0 : index + %lhs_m1 = index.add %subgroup_token_add, %c16 : index + %lhs0 = scf.if %packed_input -> (vector<16xf16>) { + %token = index.add %token_tile_origin, %lhs_m0 : index + %global_k = index.add %k_origin, %k_half : index + %loaded = func.call @ggml_packed_f16_lhs_fragment_load(%input_noalias, %bounded_token_count, %bounded_input_size, %token, %global_k) : (buffer, index, index, index, index) -> (vector<16xf16>) + scf.yield %loaded : vector<16xf16> + } else { + %loaded = vector.fragment.load %activation_stage_physical_view[%lhs_m0, %k_half] shape [%m, %k] : view<256x40xf16> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } + %lhs1 = scf.if %packed_input -> (vector<16xf16>) { + %token = index.add %token_tile_origin, %lhs_m1 : index + %global_k = index.add %k_origin, %k_half : index + %loaded = func.call @ggml_packed_f16_lhs_fragment_load(%input_noalias, %bounded_token_count, %bounded_input_size, %token, %global_k) : (buffer, index, index, index, index) -> (vector<16xf16>) + scf.yield %loaded : vector<16xf16> + } else { + %loaded = vector.fragment.load %activation_stage_physical_view[%lhs_m1, %k_half] shape [%m, %k] : view<256x40xf16> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } + %weight_fragment_layout = encoding.layout.strided [1, %weight_stride] : encoding + %weight_fragment_view = buffer.view %weight_stage[%stage_offset] : buffer -> view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> + %weight_k_group = index.mul %quant_group, %c32 : index + %weight_half = index.add %weight_k_group, %k_half : index + %rhs0 = vector.fragment.load %weight_fragment_view[%weight_half, %rhs_column0] shape [%k, %n] : view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> -> vector<16xf16> + %early_rhs1 = scf.if %serial_rhs -> (vector<16xf16>) { + scf.yield %zero_f16 : vector<16xf16> + } else { + %loaded = vector.fragment.load %weight_fragment_view[%weight_half, %rhs_column1] shape [%k, %n] : view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } + %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %rhs1 = scf.if %serial_rhs -> (vector<16xf16>) { + scf.schedule.fence + %loaded = vector.fragment.load %weight_fragment_view[%weight_half, %rhs_column1] shape [%k, %n] : view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } else { + scf.yield %early_rhs1 : vector<16xf16> + } + %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.schedule.fence + %rhs2 = vector.fragment.load %weight_fragment_view[%weight_half, %rhs_column2] shape [%k, %n] : view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> -> vector<16xf16> + %early_rhs3 = scf.if %serial_rhs -> (vector<16xf16>) { + scf.yield %zero_f16 : vector<16xf16> + } else { + %loaded = vector.fragment.load %weight_fragment_view[%weight_half, %rhs_column3] shape [%k, %n] : view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } + %half_next02 = vector.mma %lhs0, %rhs2, %half_acc02 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next12 = vector.mma %lhs1, %rhs2, %half_acc12 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %rhs3 = scf.if %serial_rhs -> (vector<16xf16>) { + scf.schedule.fence + %loaded = vector.fragment.load %weight_fragment_view[%weight_half, %rhs_column3] shape [%k, %n] : view<[%weight_window]x[%weight_rows]xf16, %weight_fragment_layout> -> vector<16xf16> + scf.yield %loaded : vector<16xf16> + } else { + scf.yield %early_rhs3 : vector<16xf16> + } + %half_next03 = vector.mma %lhs0, %rhs3, %half_acc03 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next13 = vector.mma %lhs1, %rhs3, %half_acc13 : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.schedule.fence + scf.yield %half_next00, %half_next01, %half_next02, %half_next03, %half_next10, %half_next11, %half_next12, %half_next13 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + %last_group = index.cmp eq, %quant_group, %c1 : index + %release_stage = scalar.ori %staged_input, %last_group : i1 + %release_weights = scalar.andi %release_stage, %sync_weights : i1 + scf.if %release_weights { + kernel.barrier scope(workgroup) ordering(acq_rel) + } + func.return %next00, %next01, %next02, %next03, %next10, %next11, %next12, %next13, %next_a0, %next_a1 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16> +} + + +func.def inline @ggml_f16_prefill_publish_weights(%weight_format: index, %paired: i1, %packed_input: i1, %short_weight_packets: i1, %active: i1, %bounded_input_size: index, %bounded_output_size: index, %channel_tile: index, %weight_epoch: index, %weight_noalias: buffer, %up_weight: buffer, %weight_stage: buffer, %stage_offset: offset) { + %q6_format = index.constant 6 : index + %q6_weights = index.cmp eq, %weight_format, %q6_format : index + %layout_true = scalar.constant true : i1 + %staged_input = scalar.xori %packed_input, %layout_true : i1 + %workgroup_size = kernel.workgroup.size : index + %large_workgroup = index.constant 512 : index + %wide_tile = index.cmp uge, %workgroup_size, %large_workgroup : index + %can_prefetch_a = index.cmp eq, %workgroup_size, %large_workgroup : index + %prefetch_a = scalar.andi %can_prefetch_a, %staged_input : i1 + %largest_workgroup = index.constant 1024 : index + %wg1024 = index.cmp eq, %workgroup_size, %largest_workgroup : index + %larger_tile = scalar.andi %wg1024, %staged_input : i1 + %true = scalar.constant true : i1 + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c72 = index.constant 72 : index + %c136 = index.constant 136 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %weight_stride = scf.select %wide_tile, %c72, %c136 : index + %ordinary_weight_rows = scf.select %wide_tile, %c128, %c64 : index + %staged_weight_rows = scf.select %larger_tile, %c256, %ordinary_weight_rows : index + %packed_weight_rows = scf.select %wg1024, %c128, %c64 : index + %weight_rows = scf.select %packed_input, %packed_weight_rows, %staged_weight_rows : index + %zero_weight_values = vector.constant 0.0 : vector<16xf16> + %quant_block_count = index.div %bounded_input_size, %c256 : index + %epochs_per_block = scf.select %wide_tile, %c4, %c2 : index + %weight_stage_view = buffer.view %weight_stage[%stage_offset] : buffer -> view<[%weight_rows]x[%weight_stride]xf16> + %half_weight_rows = index.div %weight_rows, %c2 : index + %tile_channels = scf.select %paired, %half_weight_rows, %weight_rows : index + %channel_tile_base = index.mul %channel_tile, %tile_channels : index + %half_weight_producers = scalar.ori %prefetch_a, %packed_input : i1 + %quant_block = index.div %weight_epoch, %epochs_per_block : index + %weight_epoch_half = index.rem %weight_epoch, %epochs_per_block : index + %groups_per_epoch = scf.select %wide_tile, %c2, %c4 : index + %weight_group_base = index.mul %weight_epoch_half, %groups_per_epoch : index + %packed_threads_per_row = scf.select %short_weight_packets, %c4, %c2 : index + %packed_producer_count = index.mul %weight_rows, %packed_threads_per_row : index + %producer_count = scf.select %packed_input, %packed_producer_count, %c256 : index + %lower_workgroup_half = index.cmp ult, %workitem, %producer_count : index + %selected_weight_producer = scf.select %half_weight_producers, %lower_workgroup_half, %true : i1 + %weight_producer = scalar.andi %selected_weight_producer, %active : i1 + scf.if %weight_producer { + %load_row_packets = scf.select %short_weight_packets, %c4, %c2 : index + %load_packet_words = scf.select %short_weight_packets, %c2, %c4 : index + %load_half = index.rem %workitem, %load_row_packets : index + %load_packet = index.mul %load_half, %load_packet_words : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %load_row_packets : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 255)] : index + %staged_threads_per_row = scf.select %half_weight_producers, %c2, %c4 : index + %weight_threads_per_row = scf.select %packed_input, %packed_threads_per_row, %staged_threads_per_row : index + %weight_local_row0 = index.div %workitem, %weight_threads_per_row : index + %weight_row_limit = index.sub %weight_rows, %c1 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, %weight_row_limit)] : index + %weight_half_row = index.rem %weight_local_row, %half_weight_rows : index + %weight_is_up = index.cmp uge, %weight_local_row, %half_weight_rows : index + %weight_pair = index.rem %load_row, %c2 : index + %weight_local_group0 = index.mul %weight_pair, %c2 : index + %weight_local_group = scf.select %wide_tile, %c0, %weight_local_group0 : index + %weight_group = index.add %weight_group_base, %weight_local_group : index + %weight_k_base = index.mul %weight_local_group, %c32 : index + %weight_k0 = index.add %weight_k_base, %load_k : index + %weight_k1 = index.add %weight_k0, %c32 : index + %channel_in_tile = scf.select %paired, %weight_half_row, %weight_local_row : index + %channel = index.add %channel_tile_base, %channel_in_tile : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { + %decoded0, %decoded1 = scf.if %q6_weights -> (vector<16xf16>, vector<16xf16>) { + %native_block_bytes = index.constant 210 : offset + %native_row_bytes = index.scale %quant_block_count, %native_block_bytes : index, offset -> offset + %row_byte_base = index.scale %channel, %native_row_bytes : index, offset -> offset + %decoded0, %decoded1 = func.call @ggml_q6k_f16_pair(%short_weight_packets, %weight_noalias, %row_byte_base, %quant_block, %weight_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> + } else { + %packed_record_bytes = index.constant 16 : offset + %packed_block_bytes = index.constant 9216 : offset + %packed_row_group = index.div %channel, %c64 : index + %packed_row_lane = index.rem %channel, %c64 : index + %packed_tile_block0 = index.mul %packed_row_group, %quant_block_count : index + %packed_tile_block = index.add %packed_tile_block0, %quant_block : index + %packed_tile_base = index.scale %packed_tile_block, %packed_block_bytes : index, offset -> offset + %packed_row_add = index.scale %packed_row_lane, %packed_record_bytes : index, offset -> offset + %row_byte_base = index.add %packed_tile_base, %packed_row_add : offset + %decoded0, %decoded1 = func.call @ggml_q4k_native_row64_f16_pair(%short_weight_packets, %paired, %weight_noalias, %up_weight, %weight_is_up, %row_byte_base, %weight_group, %load_packet) : (i1, i1, buffer, buffer, i1, offset, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> + } + scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> + } else { + scf.yield %zero_weight_values, %zero_weight_values : vector<16xf16>, vector<16xf16> + } + scf.if %short_weight_packets { + %half0 = vector.slice %weight_values0[0] : vector<16xf16> -> vector<8xf16> + %half1 = vector.slice %weight_values1[0] : vector<16xf16> -> vector<8xf16> + vector.store %half0, %weight_stage_view[%weight_local_row, %weight_k0] : vector<8xf16>, view<[%weight_rows]x[%weight_stride]xf16> + vector.store %half1, %weight_stage_view[%weight_local_row, %weight_k1] : vector<8xf16>, view<[%weight_rows]x[%weight_stride]xf16> + scf.yield + } else { + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k0] : vector<16xf16>, view<[%weight_rows]x[%weight_stride]xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k1] : vector<16xf16>, view<[%weight_rows]x[%weight_stride]xf16> + scf.yield + } + scf.yield + } + func.return +} + +func.def inline @ggml_prefill_conv4_strip(%deferred_state: i1, %width: index, %channel_base: index, %stage: buffer, %state: buffer, %filter: buffer, %output: buffer, %cache: buffer, %first: vector<8xf32>, %second: vector<8xf32>) { + %zero = index.constant 0 : index + %one = index.constant 1 : index + %two = index.constant 2 : index + %three = index.constant 3 : index + %six = index.constant 6 : index + %zero_values = vector.constant 0.0 : vector<2xf32> + %cache_rows = scf.select %deferred_state, %six, %three : index + %eight = index.constant 8 : index + %sixteen = index.constant 16 : index + %thirtytwo = index.constant 32 : index + %sixtyfour = index.constant 64 : index + %five09 = index.constant 509 : index + %five12 = index.constant 512 : index + %base = index.constant 0 : offset + %filter_bytes = index.constant 16 : offset + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 511)] : index + %wave0 = kernel.subgroup.id : index + %wave = index.assume %wave0 [range(%wave0, 0, 15)] : index + %local_row = index.mul %wave, %thirtytwo : index + %second_row = index.add %local_row, %sixteen : index + %column_lane = index.rem %tid, %eight : index + %local_column = index.mul %column_lane, %two : index + %channel0 = index.add %channel_base, %local_column : index + %channel_last = index.sub %width, %two : index + %channel = index.assume %channel0 [range(%channel0, 0, %channel_last), mul(%channel0, 2)] : index + %channel1 = index.add %channel, %one : index + %thread_row = index.div %tid, %eight : index + %tile = buffer.view %stage[%base] : buffer -> view<512x16xf32> + %output_view = buffer.view %output[%base] : buffer -> view<512x[%width]xf32> + %state_view = buffer.view %state[%base] : buffer -> view<[%width]x3xf32> + %cache_view = buffer.view %cache[%base] : buffer -> view<[%width]x[%cache_rows]xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + vector.fragment.store %first, %tile[%local_row, %zero] shape [%sixteen, %sixteen] : vector<8xf32>, view<512x16xf32> + vector.fragment.store %second, %tile[%second_row, %zero] shape [%sixteen, %sixteen] : vector<8xf32>, view<512x16xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %filter_offset = index.scale %channel, %filter_bytes : index, offset -> offset + %filter_view = buffer.view %filter[%filter_offset] : buffer -> view<2x4xf32> + %weights0 = vector.load %filter_view[%zero, %zero] : view<2x4xf32> -> vector<2x4xf32> + %weights = vector.transpose<[1, 0]> %weights0 : vector<2x4xf32> -> vector<4x2xf32> + %w0 = vector.extract %weights[0] : vector<4x2xf32> -> vector<2xf32> + %w1 = vector.extract %weights[1] : vector<4x2xf32> -> vector<2xf32> + %w2 = vector.extract %weights[2] : vector<4x2xf32> -> vector<2xf32> + %w3 = vector.extract %weights[3] : vector<4x2xf32> -> vector<2xf32> + scf.for %row_block = [%zero to %five12 step %sixtyfour] unroll { + %row = index.add %thread_row, %row_block : index + %history0 = index.cmp ult, %row, %three : index + %x0 = scf.if %history0 -> (vector<2xf32>) { + %history_values = scf.if %deferred_state -> (vector<2xf32>) { + scf.yield %zero_values : vector<2xf32> + } else { + %history_row0_0 = index.add %row, %zero : index + %history_row0 = index.assume %history_row0_0 [range(%history_row0_0, 0, 2)] : index + %h0_0 = view.load %state_view[%channel, %history_row0] : view<[%width]x3xf32> -> f32 + %h0_1 = view.load %state_view[%channel1, %history_row0] : view<[%width]x3xf32> -> f32 + %hv0_0 = vector.splat %h0_0 : vector<1xf32> + %hv0_1 = vector.splat %h0_1 : vector<1xf32> + %hv0 = vector.concat<0> %hv0_0, %hv0_1 : vector<1xf32>, vector<1xf32> -> vector<2xf32> + scf.yield %hv0 : vector<2xf32> + } + scf.yield %history_values : vector<2xf32> + } else { + %r0_0 = index.sub %row, %three : index + %r0 = index.assume %r0_0 [range(%r0_0, 0, 510)] : index + %local0 = vector.load %tile[%r0, %local_column] : view<512x16xf32> -> vector<2xf32> + scf.yield %local0 : vector<2xf32> + } + %history1 = index.cmp ult, %row, %two : index + %x1 = scf.if %history1 -> (vector<2xf32>) { + %history_values = scf.if %deferred_state -> (vector<2xf32>) { + scf.yield %zero_values : vector<2xf32> + } else { + %history_row1_0 = index.add %row, %one : index + %history_row1 = index.assume %history_row1_0 [range(%history_row1_0, 0, 2)] : index + %h1_0 = view.load %state_view[%channel, %history_row1] : view<[%width]x3xf32> -> f32 + %h1_1 = view.load %state_view[%channel1, %history_row1] : view<[%width]x3xf32> -> f32 + %hv1_0 = vector.splat %h1_0 : vector<1xf32> + %hv1_1 = vector.splat %h1_1 : vector<1xf32> + %hv1 = vector.concat<0> %hv1_0, %hv1_1 : vector<1xf32>, vector<1xf32> -> vector<2xf32> + scf.yield %hv1 : vector<2xf32> + } + scf.yield %history_values : vector<2xf32> + } else { + %r1_0 = index.sub %row, %two : index + %r1 = index.assume %r1_0 [range(%r1_0, 0, 510)] : index + %local1 = vector.load %tile[%r1, %local_column] : view<512x16xf32> -> vector<2xf32> + scf.yield %local1 : vector<2xf32> + } + %history2 = index.cmp ult, %row, %one : index + %x2 = scf.if %history2 -> (vector<2xf32>) { + %history_values = scf.if %deferred_state -> (vector<2xf32>) { + scf.yield %zero_values : vector<2xf32> + } else { + %history_row2_0 = index.add %row, %two : index + %history_row2 = index.assume %history_row2_0 [range(%history_row2_0, 0, 2)] : index + %h2_0 = view.load %state_view[%channel, %history_row2] : view<[%width]x3xf32> -> f32 + %h2_1 = view.load %state_view[%channel1, %history_row2] : view<[%width]x3xf32> -> f32 + %hv2_0 = vector.splat %h2_0 : vector<1xf32> + %hv2_1 = vector.splat %h2_1 : vector<1xf32> + %hv2 = vector.concat<0> %hv2_0, %hv2_1 : vector<1xf32>, vector<1xf32> -> vector<2xf32> + scf.yield %hv2 : vector<2xf32> + } + scf.yield %history_values : vector<2xf32> + } else { + %r2_0 = index.sub %row, %one : index + %r2 = index.assume %r2_0 [range(%r2_0, 0, 510)] : index + %local2 = vector.load %tile[%r2, %local_column] : view<512x16xf32> -> vector<2xf32> + scf.yield %local2 : vector<2xf32> + } + %x3 = vector.load %tile[%row, %local_column] : view<512x16xf32> -> vector<2xf32> + %a0 = vector.mulf %x0, %w0 : vector<2xf32> + %a1 = vector.fmaf %x1, %w1, %a0 : vector<2xf32> + %a2 = vector.fmaf %x2, %w2, %a1 : vector<2xf32> + %a3 = vector.fmaf %x3, %w3, %a2 : vector<2xf32> + %activated = vector.siluf %a3 : vector<2xf32> + vector.store %activated, %output_view[%row, %channel] : vector<2xf32>, view<512x[%width]xf32> + scf.if %deferred_state { + %first_worker = index.cmp ult, %row, %three : index + scf.if %first_worker { + %edge_row = index.assume %row [range(%row, 0, 2)] : index + %edge0 = vector.extract %x3[0] : vector<2xf32> -> f32 + %edge1 = vector.extract %x3[1] : vector<2xf32> -> f32 + view.store %edge0, %cache_view[%channel, %edge_row] : f32, view<[%width]x[%cache_rows]xf32> + view.store %edge1, %cache_view[%channel1, %edge_row] : f32, view<[%width]x[%cache_rows]xf32> + } + } + %cache_worker = index.cmp uge, %row, %five09 : index + scf.if %cache_worker { + %cache_start = scf.select %deferred_state, %three, %zero : index + %tail_row = index.sub %row, %five09 : index + %cache_row0 = index.add %tail_row, %cache_start : index + %cache_row, %cache_size = index.assume %cache_row0, %cache_rows [lt(%cache_row0, %cache_rows)] : index, index + %cache_v0 = vector.extract %x3[0] : vector<2xf32> -> f32 + %cache_v1 = vector.extract %x3[1] : vector<2xf32> -> f32 + view.store %cache_v0, %cache_view[%channel, %cache_row] : f32, view<[%width]x[%cache_rows]xf32> + view.store %cache_v1, %cache_view[%channel1, %cache_row] : f32, view<[%width]x[%cache_rows]xf32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + func.return +} + +func.def inline @ggml_mul_mat_quantized_f16_prefill_wave32_body(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer, %conv_epilogue: i1, %conv_state: buffer, %conv_filter: buffer, %conv_output: buffer, %conv_cache: buffer, %deferred_state: i1) { + %q6_format = index.constant 6 : index + %q6_weights = index.cmp eq, %weight_format, %q6_format : index + %layout_true = scalar.constant true : i1 + %unpaired = scalar.xori %paired, %layout_true : i1 + %packed_single = scalar.andi %unpaired, %packed_input : i1 + %staged_input = scalar.xori %packed_input, %layout_true : i1 + %workgroup_size = kernel.workgroup.size : index + %large_workgroup = index.constant 512 : index + %wide_tile = index.cmp uge, %workgroup_size, %large_workgroup : index + %can_prefetch_a = index.cmp eq, %workgroup_size, %large_workgroup : index + %prefetch_a = scalar.andi %can_prefetch_a, %staged_input : i1 + %largest_workgroup = index.constant 1024 : index + %wg1024 = index.cmp eq, %workgroup_size, %largest_workgroup : index + %larger_tile = scalar.andi %wg1024, %staged_input : i1 + %pipeline_weights = scalar.andi %packed_input, %can_prefetch_a : i1 + %default_short_packets = scalar.ori %larger_tile, %packed_single : i1 + %short_weight_packets = scalar.ori %default_short_packets, %pipeline_weights : i1 + %true = scalar.constant true : i1 + %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, 256)] : 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_small = index.constant 7 : index + %subgroup_large = index.constant 15 : index + %subgroup_larger = index.constant 31 : index + %ordinary_subgroup_limit = scf.select %wide_tile, %subgroup_large, %subgroup_small : index + %subgroup_limit = scf.select %wg1024, %subgroup_larger, %ordinary_subgroup_limit : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, %subgroup_limit)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c72 = index.constant 72 : index + %c136 = index.constant 136 : index + %c48 = index.constant 48 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c512 = index.constant 512 : index + %c0_offset = index.constant 0 : offset + %weight_stride = scf.select %wide_tile, %c72, %c136 : index + %ordinary_weight_rows = scf.select %wide_tile, %c128, %c64 : index + %staged_weight_rows = scf.select %larger_tile, %c256, %ordinary_weight_rows : index + %packed_weight_rows = scf.select %wg1024, %c128, %c64 : index + %weight_rows = scf.select %packed_input, %packed_weight_rows, %staged_weight_rows : index + %weight_elements = index.mul %weight_rows, %weight_stride : index + %element_bytes = index.constant 2 : offset + %weight_stage_bytes = index.scale %weight_elements, %element_bytes : index, offset -> offset + %activation_stage_bytes = index.constant 20480 : offset + %zero_weight_values = vector.constant 0.0 : vector<16xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %quant_block_count = index.div %bounded_input_size, %c256 : index + %epochs_per_block = scf.select %wide_tile, %c4, %c2 : index + %weight_epoch_count = index.mul %quant_block_count, %epochs_per_block : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %slot_count = scf.select %pipeline_weights, %c2, %c1 : index + %weight_ring_bytes = index.scale %slot_count, %weight_stage_bytes : index, offset -> offset + %conv_stage_bytes = index.constant 32768 : offset + %stage_bytes = scf.select %conv_epilogue, %conv_stage_bytes, %weight_ring_bytes : offset + %weight_stage = buffer.alloca align(16) %stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<[%weight_rows]x[%weight_stride]xf16> + %half_weight_rows = index.div %weight_rows, %c2 : index + %tile_channels = scf.select %paired, %half_weight_rows, %weight_rows : index + %channel_tile_base = index.mul %channel_tile, %tile_channels : index + %token_tile_size = scf.select %packed_input, %c512, %c256 : index + %token_tile_origin = index.mul %token_tile, %token_tile_size : index + %half_weight_producers = scalar.ori %prefetch_a, %packed_input : i1 + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init02 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init03 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init12 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init13 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %zero_prefetch = vector.constant 0.0 : vector<8xf16> + %initial_a0, %initial_a1 = scf.if %prefetch_a -> (vector<8xf16>, vector<8xf16>) { + %a0, %a1 = func.call @ggml_f16_prefill_load_pair(%input_noalias, %bounded_token_count, %bounded_input_size, %token_tile_origin, %c0) : (buffer, index, index, index, index) -> (vector<8xf16>, vector<8xf16>) + scf.yield %a0, %a1 : vector<8xf16>, vector<8xf16> + } else { + scf.yield %zero_prefetch, %zero_prefetch : vector<8xf16>, vector<8xf16> + } + scf.if %pipeline_weights { + func.call @ggml_f16_prefill_publish_weights(%weight_format, %paired, %packed_input, %short_weight_packets, %true, %bounded_input_size, %bounded_output_size, %channel_tile, %c0, %weight_noalias, %up_weight, %weight_stage, %c0_offset) : (index, i1, i1, i1, i1, index, index, index, index, buffer, buffer, buffer, offset) + kernel.barrier scope(workgroup) ordering(acq_rel) + } + %result00, %result01, %result02, %result03, %result10, %result11, %result12, %result13, %final_a0, %final_a1 = scf.for %weight_epoch = [%c0 to %weight_epoch_count step %c1](%block_acc00 = %init00 : vector<8xf32>, %block_acc01 = %init01 : vector<8xf32>, %block_acc02 = %init02 : vector<8xf32>, %block_acc03 = %init03 : vector<8xf32>, %block_acc10 = %init10 : vector<8xf32>, %block_acc11 = %init11 : vector<8xf32>, %block_acc12 = %init12 : vector<8xf32>, %block_acc13 = %init13 : vector<8xf32>, %block_a0 = %initial_a0 : vector<8xf16>, %block_a1 = %initial_a1 : vector<8xf16>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) { + %quant_block = index.div %weight_epoch, %epochs_per_block : index + %weight_epoch_half = index.rem %weight_epoch, %epochs_per_block : index + %groups_per_epoch = scf.select %wide_tile, %c2, %c4 : index + %weight_group_base = index.mul %weight_epoch_half, %groups_per_epoch : index + %slot = index.rem %weight_epoch, %c2 : index + %slot_offset = index.scale %slot, %weight_stage_bytes : index, offset -> offset + %stage_offset = scf.select %pipeline_weights, %slot_offset, %c0_offset : offset + %next_epoch = index.add %weight_epoch, %c1 : index + %load_epoch = scf.select %pipeline_weights, %next_epoch, %weight_epoch : index + %has_load = index.cmp ult, %load_epoch, %weight_epoch_count : index + %load_slot = index.rem %load_epoch, %c2 : index + %load_slot_offset = index.scale %load_slot, %weight_stage_bytes : index, offset -> offset + %load_offset = scf.select %pipeline_weights, %load_slot_offset, %c0_offset : offset + func.call @ggml_f16_prefill_publish_weights(%weight_format, %paired, %packed_input, %short_weight_packets, %has_load, %bounded_input_size, %bounded_output_size, %channel_tile, %load_epoch, %weight_noalias, %up_weight, %weight_stage, %load_offset) : (index, i1, i1, i1, i1, index, index, index, index, buffer, buffer, buffer, offset) + %not_wide_tile = scalar.xori %wide_tile, %true : i1 + // First activation publication also publishes the wide weight tile. + scf.if %not_wide_tile { + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield + } + %block_result00, %block_result01, %block_result02, %block_result03, %block_result10, %block_result11, %block_result12, %block_result13, %block_next_a0, %block_next_a1 = scf.if %wide_tile -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) { + %first00, %first01, %first02, %first03, %first10, %first11, %first12, %first13, %first_a0, %first_a1 = func.call @ggml_f16_prefill_stage_mma(%paired, %packed_input, %bounded_token_count, %bounded_input_size, %token_tile_origin, %quant_block, %weight_group_base, %c0, %input_noalias, %weight_stage, %activation_stage, %stage_offset, %pipeline_weights, %block_acc00, %block_acc01, %block_acc02, %block_acc03, %block_acc10, %block_acc11, %block_acc12, %block_acc13, %block_a0, %block_a1) : (i1, i1, index, index, index, index, index, index, buffer, buffer, buffer, offset, i1, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) + %second00, %second01, %second02, %second03, %second10, %second11, %second12, %second13, %second_a0, %second_a1 = func.call @ggml_f16_prefill_stage_mma(%paired, %packed_input, %bounded_token_count, %bounded_input_size, %token_tile_origin, %quant_block, %weight_group_base, %c1, %input_noalias, %weight_stage, %activation_stage, %stage_offset, %pipeline_weights, %first00, %first01, %first02, %first03, %first10, %first11, %first12, %first13, %first_a0, %first_a1) : (i1, i1, index, index, index, index, index, index, buffer, buffer, buffer, offset, i1, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) + scf.yield %second00, %second01, %second02, %second03, %second10, %second11, %second12, %second13, %second_a0, %second_a1 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16> + } else { + %ordinary00, %ordinary01, %ordinary02, %ordinary03, %ordinary10, %ordinary11, %ordinary12, %ordinary13 = scf.for %quant_group = [%c0 to %c4 step %c1](%acc00 = %block_acc00 : vector<8xf32>, %acc01 = %block_acc01 : vector<8xf32>, %acc02 = %block_acc02 : vector<8xf32>, %acc03 = %block_acc03 : vector<8xf32>, %acc10 = %block_acc10 : vector<8xf32>, %acc11 = %block_acc11 : vector<8xf32>, %acc12 = %block_acc12 : vector<8xf32>, %acc13 = %block_acc13 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %next00, %next01, %next02, %next03, %next10, %next11, %next12, %next13, %unused_a0, %unused_a1 = func.call @ggml_f16_prefill_stage_mma(%paired, %packed_input, %bounded_token_count, %bounded_input_size, %token_tile_origin, %quant_block, %weight_group_base, %quant_group, %input_noalias, %weight_stage, %activation_stage, %stage_offset, %pipeline_weights, %acc00, %acc01, %acc02, %acc03, %acc10, %acc11, %acc12, %acc13, %zero_prefetch, %zero_prefetch) : (i1, i1, index, index, index, index, index, index, buffer, buffer, buffer, offset, i1, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16>) + scf.yield %next00, %next01, %next02, %next03, %next10, %next11, %next12, %next13 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + scf.yield %ordinary00, %ordinary01, %ordinary02, %ordinary03, %ordinary10, %ordinary11, %ordinary12, %ordinary13, %zero_prefetch, %zero_prefetch : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16> + } + scf.if %pipeline_weights { + kernel.barrier scope(workgroup) ordering(acq_rel) + } + scf.yield %block_result00, %block_result01, %block_result02, %block_result03, %block_result10, %block_result11, %block_result12, %block_result13, %block_next_a0, %block_next_a1 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf16>, vector<8xf16> + } + %token_subgroups = scf.select %packed_input, %c16, %c8 : index + %token_subgroup = index.rem %subgroup, %token_subgroups : index + %subgroup_token_add = index.mul %token_subgroup, %c32 : index + %token_tile_base = index.add %token_tile_origin, %subgroup_token_add : index + %channel_subgroup = index.div %subgroup, %token_subgroups : index + %subgroup_channel_stride = scf.select %paired, %c32, %c64 : index + %subgroup_channel_add = index.mul %channel_subgroup, %subgroup_channel_stride : index + %channel_output_base = index.add %channel_tile_base, %subgroup_channel_add : index + scf.if %conv_epilogue { + %conv_channel_0 = index.add %channel_output_base, %c0 : index + func.call @ggml_prefill_conv4_strip(%deferred_state, %bounded_output_size, %conv_channel_0, %weight_stage, %conv_state, %conv_filter, %conv_output, %conv_cache, %result00, %result10) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, vector<8xf32>, vector<8xf32>) + %conv_channel_1 = index.add %channel_output_base, %c16 : index + func.call @ggml_prefill_conv4_strip(%deferred_state, %bounded_output_size, %conv_channel_1, %weight_stage, %conv_state, %conv_filter, %conv_output, %conv_cache, %result01, %result11) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, vector<8xf32>, vector<8xf32>) + %conv_channel_2 = index.add %channel_output_base, %c32 : index + func.call @ggml_prefill_conv4_strip(%deferred_state, %bounded_output_size, %conv_channel_2, %weight_stage, %conv_state, %conv_filter, %conv_output, %conv_cache, %result02, %result12) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, vector<8xf32>, vector<8xf32>) + %conv_channel_3 = index.add %channel_output_base, %c48 : index + func.call @ggml_prefill_conv4_strip(%deferred_state, %bounded_output_size, %conv_channel_3, %weight_stage, %conv_state, %conv_filter, %conv_output, %conv_cache, %result03, %result13) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, vector<8xf32>, vector<8xf32>) + } else { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + %f16_columns = scf.select %packed_output, %c16, %bounded_output_size : index + %f16_group_elements = index.mul %channel_output_base, %bounded_token_count : index + %f16_group_bytes = index.scale %f16_group_elements, %element_bytes : index, offset -> offset + %f16_next_elements = index.mul %c16, %bounded_token_count : index + %f16_next_bytes = index.scale %f16_next_elements, %element_bytes : index, offset -> offset + %f16_next_group_bytes = index.add %f16_group_bytes, %f16_next_bytes : offset + %f16_base0 = scf.select %packed_output, %f16_group_bytes, %c0_offset : offset + %f16_base1 = scf.select %packed_output, %f16_next_group_bytes, %c0_offset : offset + %f16_output_view0 = buffer.view %f16_output[%f16_base0] : buffer -> view<[%bounded_token_count]x[%f16_columns]xf16> + %f16_output_view1 = buffer.view %f16_output[%f16_base1] : buffer -> view<[%bounded_token_count]x[%f16_columns]xf16> + scf.if %paired { + %out_m00 = index.add %token_tile_base, %c0 : index + %out_n00 = index.add %channel_output_base, %c0 : index + %fused00 = template.apply<@ggml.binary_f32.apply_vector8>(%binary_op, %result00, %result02) : (index, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + vector.fragment.store %fused00, %output_view[%out_m00, %out_n00] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %narrow00 = vector.fptrunc %fused00 : vector<8xf32> to vector<8xf16> + %f16_n00 = scf.select %packed_output, %c0, %out_n00 : index + vector.fragment.store %narrow00, %f16_output_view0[%out_m00, %f16_n00] shape [%m, %n] : vector<8xf16>, view<[%bounded_token_count]x[%f16_columns]xf16> + scf.schedule.fence + %out_m01 = index.add %token_tile_base, %c0 : index + %out_n01 = index.add %channel_output_base, %c16 : index + %fused01 = template.apply<@ggml.binary_f32.apply_vector8>(%binary_op, %result01, %result03) : (index, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + vector.fragment.store %fused01, %output_view[%out_m01, %out_n01] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %narrow01 = vector.fptrunc %fused01 : vector<8xf32> to vector<8xf16> + %f16_n01 = scf.select %packed_output, %c0, %out_n01 : index + vector.fragment.store %narrow01, %f16_output_view1[%out_m01, %f16_n01] shape [%m, %n] : vector<8xf16>, view<[%bounded_token_count]x[%f16_columns]xf16> + scf.schedule.fence + %out_m10 = index.add %token_tile_base, %c16 : index + %out_n10 = index.add %channel_output_base, %c0 : index + %fused10 = template.apply<@ggml.binary_f32.apply_vector8>(%binary_op, %result10, %result12) : (index, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + vector.fragment.store %fused10, %output_view[%out_m10, %out_n10] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %narrow10 = vector.fptrunc %fused10 : vector<8xf32> to vector<8xf16> + %f16_n10 = scf.select %packed_output, %c0, %out_n10 : index + vector.fragment.store %narrow10, %f16_output_view0[%out_m10, %f16_n10] shape [%m, %n] : vector<8xf16>, view<[%bounded_token_count]x[%f16_columns]xf16> + scf.schedule.fence + %out_m11 = index.add %token_tile_base, %c16 : index + %out_n11 = index.add %channel_output_base, %c16 : index + %fused11 = template.apply<@ggml.binary_f32.apply_vector8>(%binary_op, %result11, %result13) : (index, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + vector.fragment.store %fused11, %output_view[%out_m11, %out_n11] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %narrow11 = vector.fptrunc %fused11 : vector<8xf32> to vector<8xf16> + %f16_n11 = scf.select %packed_output, %c0, %out_n11 : index + vector.fragment.store %narrow11, %f16_output_view1[%out_m11, %f16_n11] shape [%m, %n] : vector<8xf16>, view<[%bounded_token_count]x[%f16_columns]xf16> + scf.schedule.fence + } else { + %out_m00 = index.add %token_tile_base, %c0 : index + %out_n00 = index.add %channel_output_base, %c0 : index + vector.fragment.store %result00, %output_view[%out_m00, %out_n00] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m01 = index.add %token_tile_base, %c0 : index + %out_n01 = index.add %channel_output_base, %c16 : index + vector.fragment.store %result01, %output_view[%out_m01, %out_n01] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m02 = index.add %token_tile_base, %c0 : index + %out_n02 = index.add %channel_output_base, %c32 : index + vector.fragment.store %result02, %output_view[%out_m02, %out_n02] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m03 = index.add %token_tile_base, %c0 : index + %out_n03 = index.add %channel_output_base, %c48 : index + vector.fragment.store %result03, %output_view[%out_m03, %out_n03] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m10 = index.add %token_tile_base, %c16 : index + %out_n10 = index.add %channel_output_base, %c0 : index + vector.fragment.store %result10, %output_view[%out_m10, %out_n10] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m11 = index.add %token_tile_base, %c16 : index + %out_n11 = index.add %channel_output_base, %c16 : index + vector.fragment.store %result11, %output_view[%out_m11, %out_n11] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m12 = index.add %token_tile_base, %c16 : index + %out_n12 = index.add %channel_output_base, %c32 : index + vector.fragment.store %result12, %output_view[%out_m12, %out_n12] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + %out_m13 = index.add %token_tile_base, %c16 : index + %out_n13 = index.add %channel_output_base, %c48 : index + vector.fragment.store %result13, %output_view[%out_m13, %out_n13] shape [%m, %n] : vector<8xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32> + } + } + func.return +} + +func.def public inline @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer) { + %no_conv = scalar.constant false : i1 + func.call @ggml_mul_mat_quantized_f16_prefill_wave32_body(%weight_format, %binary_op, %paired, %packed_input, %packed_output, %token_count, %input_size0, %output_size0, %channel_tile, %token_tile, %input, %weight, %up_weight, %output, %f16_output, %no_conv, %output, %output, %output, %output, %no_conv) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, i1, buffer, buffer, buffer, buffer, i1) + func.return +} + +func.def public inline @ggml_mul_mat_quantized_f16_prefill_conv4(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %input: buffer, %weight: buffer, %state: buffer, %filter: buffer, %output: buffer, %cache: buffer) { + %zero = index.constant 0 : index + %false = scalar.constant false : i1 + %true = scalar.constant true : i1 + func.call @ggml_mul_mat_quantized_f16_prefill_wave32_body(%weight_format, %zero, %false, %true, %false, %token_count, %input_size, %output_size, %channel_tile, %zero, %input, %weight, %weight, %output, %output, %true, %state, %filter, %output, %cache, %false) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, i1, buffer, buffer, buffer, buffer, i1) + func.return +} + +func.def public inline @ggml_mul_mat_quantized_f16_prefill_conv4_interior(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %input: buffer, %weight: buffer, %filter: buffer, %output: buffer, %edges: buffer) { + %zero = index.constant 0 : index + %false = scalar.constant false : i1 + %true = scalar.constant true : i1 + func.call @ggml_mul_mat_quantized_f16_prefill_wave32_body(%weight_format, %zero, %false, %true, %false, %token_count, %input_size, %output_size, %channel_tile, %zero, %input, %weight, %weight, %output, %output, %true, %edges, %filter, %output, %edges, %true) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, i1, buffer, buffer, buffer, buffer, i1) + func.return +} + +// Exact K16-major F16 activation tiles used by larger token-reuse GEMMs. +func.def inline @ggml_packed_f16_lhs_fragment_load(%input: buffer, %tokens: index, %input_size: index, %token: index, %k_origin: index) -> (vector<16xf16>) { + %zero = index.constant 0 : offset + %c0 = index.constant 0 : index + %c16 = index.constant 16 : index + %k_tiles = index.div %input_size, %c16 : index + %k_tile = index.div %k_origin, %c16 : index + %packed_row = index.madd %k_tile, %tokens, %token : index + %packed_rows = index.mul %tokens, %k_tiles : index + %packed_view = buffer.view %input[%zero] : buffer -> view<[%packed_rows]x16xf16> + %values = vector.fragment.load %packed_view[%packed_row, %c0] shape [%c16, %c16] : view<[%packed_rows]x16xf16> -> vector<16xf16> + func.return %values : vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/publish_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/publish_f16.loom new file mode 100644 index 000000000000..018e56dbed1a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/publish_f16.loom @@ -0,0 +1,49 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.normalization.publish_f16.k16_vector4(%values: vector<4xf32>, %channel: index, %token: index, %token_count: index, %hidden_size: index, %packet: index, %output: buffer) + +template.def<@ggml.normalization.publish_f16.k16_vector4> device @ggml_normalization_publish_f16_k16_vector4(%values: vector<4xf32>, %channel: index, %token: index, %token_count: index, %hidden_size: index, %packet: index, %output: buffer) { + %base = index.constant 0 : offset + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c16 = index.constant 16 : index + %packed_columns = index.div %hidden_size, %c16 : index + %packed_view = buffer.view %output[%base] : buffer -> view<[%packed_columns]x[%token_count]x16xf16> + %packed_column = index.div %channel, %c16 : index + %packed_lane = index.rem %channel, %c16 : index + %half = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store %half, %packed_view[%packed_column, %token, %packed_lane] : vector<4xf16>, view<[%packed_columns]x[%token_count]x16xf16> + %fence_pos = index.rem %packet, %c4 : index + %end_packet = index.cmp eq, %fence_pos, %c3 : index + scf.if %end_packet { + scf.schedule.fence + } + template.return +} + +template.decl @ggml.normalization.publish_f16.gate_vector4(%values: vector<4xf32>, %row: index, %channel: index, %token_count: index, %hidden_size: index, %row_width: index, %output: buffer) + +template.def<@ggml.normalization.publish_f16.gate_vector4> device @ggml_normalization_publish_f16_gate_vector4(%values: vector<4xf32>, %row: index, %channel: index, %token_count: index, %hidden_size: index, %row_width: index, %output: buffer) { + %z = index.constant 0 : index + %zero_offset = index.constant 0 : offset + %c16 = index.constant 16 : index + %output_view = buffer.view %output[%zero_offset] : buffer -> view<[%token_count]x[%hidden_size]xf16> + %half = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + %packed = index.cmp ugt, %row_width, %z : index + scf.if %packed { + %element_count = index.mul %token_count, %hidden_size : index + %packed_rows = index.div %element_count, %row_width : index + %packed_columns = index.div %row_width, %c16 : index + %packed_view = buffer.view %output[%zero_offset] : buffer -> view<[%packed_columns]x[%packed_rows]x16xf16> + %flat = index.madd %row, %hidden_size, %channel : index + %packed_row = index.div %flat, %row_width : index + %packed_channel = index.rem %flat, %row_width : index + %packed_column = index.div %packed_channel, %c16 : index + %packed_lane = index.rem %packed_channel, %c16 : index + vector.store %half, %packed_view[%packed_column, %packed_row, %packed_lane] : vector<4xf16>, view<[%packed_columns]x[%packed_rows]x16xf16> + } else { + vector.store %half, %output_view[%row, %channel] : vector<4xf16>, view<[%token_count]x[%hidden_size]xf16> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/publish_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/publish_q8.loom new file mode 100644 index 000000000000..e3f4e866d40f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/publish_q8.loom @@ -0,0 +1,26 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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) +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.normalization.publish_q8.binary_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.def<@ggml.normalization.publish_q8.binary_vector4> device @ggml_normalization_publish_q8_binary_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} + +template.decl @ggml.normalization.publish_q8.gate_vector4(%publish_word: i1, %row: index, %hidden_size: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.def<@ggml.normalization.publish_q8.gate_vector4> device @ggml_normalization_publish_q8_gate_vector4(%publish_word: i1, %row: index, %hidden_size: index, %channel: index, %values: vector<4xf32>, %output: buffer) { + %zero_offset = index.constant 0 : offset + %scratch_bytes = index.constant 1152 : offset + %scratch_d_offset = index.constant 1024 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%zero_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_offset] : buffer -> view<32xf32> + %flat_channel = index.madd %row, %hidden_size, %channel : index + template.apply<@ggml.quantize_q8_1_x4.publish_vector4_strict>(%publish_word, %zero_offset, %flat_channel, %values, %scratch_values, %scratch_d, %output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom new file mode 100644 index 000000000000..2c75f94808c7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm.loom @@ -0,0 +1,367 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Computes the reciprocal RMS scale for one F32 row. The caller owns launch +// geometry; this motif assumes a 256-workitem wave32 workgroup. +template.decl @ggml.rmsnorm_f32.apply(%value: f32, %scale: f32) -> (f32) + +template.decl @ggml.rmsnorm_f32.row_scale(%token_count0: index, %token0: index, %hidden_size0: index, %input_stride0: index, %epsilon: f32, %input: buffer) -> (f32, index, index, i1) + +template.def<@ggml.rmsnorm_f32.row_scale> device @ggml_rmsnorm_f32_row_scale(%token_count0: index, %token0: index, %hidden_size0: index, %input_stride0: index, %epsilon: f32, %input: buffer) -> (f32, index, index, i1) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1048576)] : index + %hidden_size1 = index.assume %hidden_size0 [range(%hidden_size0, 64, 32768), mul(%hidden_size0, 64)] : index + %input_stride1 = index.assume %input_stride0 [range(%input_stride0, 64, 1048576)] : index + %hidden_size, %input_stride = index.assume %hidden_size1, %input_stride1 [le(%hidden_size1, %input_stride1)] : index, index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %scratch_bytes = index.constant 1024 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %valid_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token = index.assume %safe_token0 [lt(%safe_token0, %token_count)] : index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%input_stride]xf32> + %c1024 = index.constant 1024 : index + %tail = index.rem %hidden_size, %c1024 : index + %unroll_limit = index.sub %hidden_size, %tail : index + %prefix_sum = scf.for %channel = [%workitem to %unroll_limit step %c1024](%running_sum = %c0_f32 : f32) -> (f32) { + %offset0 = index.constant 0 : index + %channel0 = index.add %channel, %offset0 : index + %value0 = view.load %input_view[%token, %channel0] : view<[%token_count]x[%input_stride]xf32> -> f32 + %offset1 = index.constant 256 : index + %channel1 = index.add %channel, %offset1 : index + %value1 = view.load %input_view[%token, %channel1] : view<[%token_count]x[%input_stride]xf32> -> f32 + %offset2 = index.constant 512 : index + %channel2 = index.add %channel, %offset2 : index + %value2 = view.load %input_view[%token, %channel2] : view<[%token_count]x[%input_stride]xf32> -> f32 + %offset3 = index.constant 768 : index + %channel3 = index.add %channel, %offset3 : index + %value3 = view.load %input_view[%token, %channel3] : view<[%token_count]x[%input_stride]xf32> -> f32 + %square0 = scalar.mulf %value0, %value0 : f32 + %sum0 = scalar.addf %running_sum, %square0 : f32 + %square1 = scalar.mulf %value1, %value1 : f32 + %sum1 = scalar.addf %sum0, %square1 : f32 + %square2 = scalar.mulf %value2, %value2 : f32 + %sum2 = scalar.addf %sum1, %square2 : f32 + %square3 = scalar.mulf %value3, %value3 : f32 + %sum3 = scalar.addf %sum2, %square3 : f32 + scf.yield %sum3 : f32 + } + %tail_start = index.add %unroll_limit, %workitem : index + %thread_sum = scf.for %channel = [%tail_start to %hidden_size step %c256](%running_sum = %prefix_sum : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%token_count]x[%input_stride]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + template.return %scale, %token, %token_count, %valid_token : f32, index, index, i1 +} + +// One subgroup owns a short row; callers supply a valid row index and pack +// independent rows into a workgroup. No workgroup scratch or barriers. +template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) + +template.def<@ggml.rmsnorm_f32.subgroup_row_scale> device @ggml_rmsnorm_f32_subgroup_row_scale(%token_count0: index, %token0: index, %hidden0: index, %epsilon: f32, %input_n: buffer) -> (f32) { + %bounded_token_count = index.assume %token_count0 [range(%token_count0, 1, 1048576)] : index + %row = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 1024), mul(%hidden0, 128)] : index + %lane = kernel.subgroup.lane.id : index + %four = index.constant 4 : index + %c128 = index.constant 128 : index + %zero_offset = index.constant 0 : offset + %channel0 = index.mul %lane, %four : index + %zero = scalar.constant 0.0 : f32 + %zero4 = vector.constant 0.0 : vector<4xf32> + %iv = buffer.view %input_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %last_channel = index.sub %hidden, %four : index + %squares = scf.for %channel = [%channel0 to %hidden step %c128](%sum = %zero4 : vector<4xf32>) -> (vector<4xf32>) { + %bounded_channel = index.assume %channel [range(%channel, 0, 1020), mul(%channel, 4), le(%channel, %last_channel)] : index + %x = vector.load %iv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> + %xx = vector.mulf %x, %x : vector<4xf32> + %next = vector.addf %sum, %xx : vector<4xf32> + scf.yield %next : vector<4xf32> + } + %lane_sum = vector.reduce %squares, %zero : vector<4xf32>, f32 + %row_sum = kernel.subgroup.reduce %lane_sum : f32 + %hidden_i = index.cast %hidden : index to i32 + %hidden_f = scalar.sitofp %hidden_i : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_f : f32 + %biased = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased : f32 + template.return %scale : f32 +} + +template.def<@ggml.rmsnorm_f32.apply> device @ggml_rmsnorm_f32_apply(%value: f32, %scale: f32) -> (f32) { + %normalized = scalar.mulf %value, %scale : f32 + template.return %normalized : f32 +} + + +config.decl @ggml.rmsnorm_f32.hidden_size : %value: index where [range(%value, 64, 32768), mul(%value, 64)] +config.decl @ggml.rmsnorm_f32.input_stride : %value: index where [range(%value, 64, 1048576)] +config.decl @ggml.rmsnorm_f32.rms_epsilon : f32 + +template.decl @ggml.rmsnorm_f32.body(%token_count: index, %input: buffer, %output: buffer) + +template.def<@ggml.rmsnorm_f32.body> device @ggml_rmsnorm_f32_body(%token_count: index, %input: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_f32.hidden_size : index + %input_stride0 = config.get @ggml.rmsnorm_f32.input_stride : index + %epsilon = config.get @ggml.rmsnorm_f32.rms_epsilon : f32 + %hidden_size1 = index.assume %hidden_size0 [range(%hidden_size0, 64, 32768), mul(%hidden_size0, 64)] : index + %input_stride1 = index.assume %input_stride0 [range(%input_stride0, 64, 1048576)] : index + %hidden_size, %input_stride = index.assume %hidden_size1, %input_stride1 [le(%hidden_size1, %input_stride1)] : index, index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %input_stride, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%input_stride]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%input_stride]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + view.store %normalized, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + template.return +} + +template.decl @ggml.rope_f32.rotate_scalar_pair(%x: f32, %y: f32, %cosine: f32, %sine: f32) -> (f32, f32) +template.decl @ggml.rmsnorm_mul_rope_f32.body(%input: buffer, %weight: buffer, %positions: buffer, %output: buffer) + +config.decl @ggml.rmsnorm_mul_rope.hidden_size : %value: index where [range(%value, 2, 512)] +config.decl @ggml.rmsnorm_mul_rope.ne1 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.rmsnorm_mul_rope.ne2 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.rmsnorm_mul_rope.ne3 : %value: index where [range(%value, 1, 1024)] +config.decl @ggml.rmsnorm_mul_rope.input_stride1 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.input_stride2 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.input_stride3 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.output_stride1 : %value: index where [range(%value, 0, 268435456)] +config.decl @ggml.rmsnorm_mul_rope.output_stride2 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.output_stride3 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.n_dims : %value: index where [range(%value, 2, 512)] +config.decl @ggml.rmsnorm_mul_rope.section0 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.section1 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.section2 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.section3 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.mode : %value: index where [range(%value, 2, 40)] +config.decl @ggml.rmsnorm_mul_rope.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] +config.decl @ggml.rmsnorm_mul_rope.epsilon : f32 +config.decl @ggml.rmsnorm_mul_rope.freq_base : f32 +config.decl @ggml.rmsnorm_mul_rope.freq_scale : f32 +config.decl @ggml.rmsnorm_mul_rope.attn_factor : f32 + +template.def<@ggml.rmsnorm_mul_rope_f32.body> device @ggml_rmsnorm_mul_rope_f32_body(%src0: buffer, %weight: buffer, %pos: buffer, %dst: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %cneox = index.constant 2 : index + %c40 = index.constant 40 : index + %c0_f32 = scalar.constant 0.0 : f32 + %cneg2_f32 = scalar.constant -2.0 : f32 + %eps = config.get @ggml.rmsnorm_mul_rope.epsilon : f32 + %freq_base = config.get @ggml.rmsnorm_mul_rope.freq_base : f32 + %freq_scale = config.get @ggml.rmsnorm_mul_rope.freq_scale : f32 + %attn_factor = config.get @ggml.rmsnorm_mul_rope.attn_factor : f32 + + %ncols = config.get @ggml.rmsnorm_mul_rope.hidden_size : index + %ne1 = config.get @ggml.rmsnorm_mul_rope.ne1 : index + %ne2 = config.get @ggml.rmsnorm_mul_rope.ne2 : index + %s1 = config.get @ggml.rmsnorm_mul_rope.input_stride1 : index + %s2 = config.get @ggml.rmsnorm_mul_rope.input_stride2 : index + %s3 = config.get @ggml.rmsnorm_mul_rope.input_stride3 : index + %d1 = config.get @ggml.rmsnorm_mul_rope.output_stride1 : index + %d2 = config.get @ggml.rmsnorm_mul_rope.output_stride2 : index + %d3 = config.get @ggml.rmsnorm_mul_rope.output_stride3 : index + %n_dims = config.get @ggml.rmsnorm_mul_rope.n_dims : index + %sec0 = config.get @ggml.rmsnorm_mul_rope.section0 : index + %sec1 = config.get @ggml.rmsnorm_mul_rope.section1 : index + %sec2 = config.get @ggml.rmsnorm_mul_rope.section2 : index + %sec3 = config.get @ggml.rmsnorm_mul_rope.section3 : index + %mode = config.get @ggml.rmsnorm_mul_rope.mode : index + %wg = config.get @ggml.rmsnorm_mul_rope.workgroup_size : index + + %ncols_b = index.assume %ncols [range(%ncols, 2, 512)] : index + %ne1_b = index.assume %ne1 [range(%ne1, 1, 1048576)] : index + %ne2_b = index.assume %ne2 [range(%ne2, 1, 1048576)] : index + %n_dims_b = index.assume %n_dims [range(%n_dims, 2, 512)] : index + + %row0 = kernel.workgroup.id : index + %row = index.assume %row0 [range(%row0, 0, 1073741823)] : index + %lane0 = kernel.workitem.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 1023)] : index + + %src_global = buffer.assume.memory_space %src0 : buffer + %weight_global = buffer.assume.memory_space %weight : buffer + %pos_global = buffer.assume.memory_space %pos : buffer + %dst_global = buffer.assume.memory_space %dst : buffer + %src_na, %weight_na, %pos_na, %dst_na = buffer.assume.noalias %src_global, %weight_global, %pos_global, %dst_global : buffer, buffer, buffer, buffer + %src_view = buffer.view %src_na[%base] : buffer -> view<1073741824xf32> + %weight_view = buffer.view %weight_na[%base] : buffer -> view<65536xf32> + %pos_view = buffer.view %pos_na[%base] : buffer -> view<1073741824xi32> + %dst_view = buffer.view %dst_na[%base] : buffer -> view<1073741824xf32> + + %n12 = index.mul %ne1_b, %ne2_b : index + %i3 = index.div %row, %n12 : index + %r12 = index.rem %row, %n12 : index + %i2 = index.div %r12, %ne1_b : index + %i1 = index.rem %r12, %ne1_b : index + %o1 = index.mul %i1, %s1 : index + %o2 = index.mul %i2, %s2 : index + %o3 = index.mul %i3, %s3 : index + %o12 = index.add %o1, %o2 : index + %src_base0 = index.add %o12, %o3 : index + %src_base = index.assume %src_base0 [range(%src_base0, 0, 1073741823)] : index + + %sum = scf.for %col = [%lane to %ncols_b step %wg](%acc = %c0_f32 : f32) -> (f32) { + %si0 = index.add %src_base, %col : index + %si = index.assume %si0 [range(%si0, 0, 1073741823)] : index + %value = view.load %src_view[%si] : view<1073741824xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next = scalar.addf %acc, %square : f32 + scf.yield %next : f32 + } + %row_sum = kernel.workgroup.reduce %sum : f32 + %ncols_i32 = index.cast %ncols_b : index to i32 + %ncols_f32 = scalar.sitofp %ncols_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %ncols_f32 : f32 + %biased = scalar.addf %mean, %eps : f32 + %rms_scale = scalar.rsqrtf %biased : f32 + + %half = index.div %ncols_b, %c2 : index + %active = index.cmp ult, %lane, %half : index + scf.if %active { + %ih = index.add %lane, %c0 : index + %i0 = index.mul %ih, %c2 : index + %half_dims = index.div %n_dims_b, %c2 : index + %rotate = index.cmp ult, %i0, %n_dims_b : index + + %sd01 = index.add %sec0, %sec1 : index + %sd012 = index.add %sd01, %sec2 : index + %sect_dims = index.add %sd012, %sec3 : index + %sector = index.rem %ih, %sect_dims : index + %m3 = index.rem %sector, %c3 : index + %s0lim = index.mul %sec0, %c3 : index + %s1lim = index.mul %sec1, %c3 : index + %s2lim = index.mul %sec2, %c3 : index + %is0 = index.cmp eq, %m3, %c0 : index + %is1 = index.cmp eq, %m3, %c1 : index + %is2 = index.cmp eq, %m3, %c2 : index + %ok0 = index.cmp ult, %sector, %s0lim : index + %ok1 = index.cmp ult, %sector, %s1lim : index + %ok2 = index.cmp ult, %sector, %s2lim : index + %im_p0 = scf.select %ok0, %c0, %c3 : index + %im_a = scf.select %is0, %im_p0, %c3 : index + %im_p2 = scf.select %ok2, %c2, %c3 : index + %im_b = scf.select %is2, %im_p2, %im_a : index + %im_p1 = scf.select %ok1, %c1, %c3 : index + %plane_im = scf.select %is1, %im_p1, %im_b : index + + %sec_w = index.add %sec0, %sec1 : index + %sec_w2 = index.add %sec_w, %sec2 : index + %lt0 = index.cmp ult, %sector, %sec0 : index + %ltw = index.cmp ult, %sector, %sec_w : index + %ltw2 = index.cmp ult, %sector, %sec_w2 : index + %mr_a = scf.select %ltw2, %c2, %c3 : index + %mr_b = scf.select %ltw, %c1, %mr_a : index + %plane_mr = scf.select %lt0, %c0, %mr_b : index + %is_im = index.cmp eq, %mode, %c40 : index + %plane = scf.select %is_im, %plane_im, %plane_mr : index + + %plane_off = index.mul %plane, %ne2_b : index + %sectioned_pos_i = index.add %i2, %plane_off : index + %is_neox = index.cmp eq, %mode, %cneox : index + %pos_i0 = scf.select %is_neox, %i2, %sectioned_pos_i : index + %pos_i = index.assume %pos_i0 [range(%pos_i0, 0, 1073741823)] : index + %pos_i32 = view.load %pos_view[%pos_i] : view<1073741824xi32> -> i32 + %pos_f = scalar.sitofp %pos_i32 : i32 to f32 + %n_dims_i32 = index.cast %n_dims_b : index to i32 + %n_dims_f32 = scalar.sitofp %n_dims_i32 : i32 to f32 + %ih_i32 = index.cast %ih : index to i32 + %ih_f = scalar.sitofp %ih_i32 : i32 to f32 + %scaled_ih = scalar.mulf %cneg2_f32, %ih_f : f32 + %theta_exponent = scalar.divf %scaled_ih, %n_dims_f32 : f32 + %freq = scalar.powf %freq_base, %theta_exponent : f32 + %theta_base = scalar.mulf %pos_f, %freq : f32 + %theta = scalar.mulf %theta_base, %freq_scale : f32 + %cos_theta_raw = scalar.cosf %theta : f32 + %s0v = scalar.sinf %theta : f32 + %cos_theta = scalar.mulf %cos_theta_raw, %attn_factor : f32 + %sin_theta = scalar.mulf %s0v, %attn_factor : f32 + + %off_a = scf.select %rotate, %ih, %i0 : index + %rot_b = index.add %ih, %half_dims : index + %pass_b = index.add %i0, %c1 : index + %off_b = scf.select %rotate, %rot_b, %pass_b : index + %sa0 = index.add %src_base, %off_a : index + %sa = index.assume %sa0 [range(%sa0, 0, 1073741823)] : index + %sb0 = index.add %src_base, %off_b : index + %sb = index.assume %sb0 [range(%sb0, 0, 1073741823)] : index + + %raw0 = view.load %src_view[%sa] : view<1073741824xf32> -> f32 + %raw1 = view.load %src_view[%sb] : view<1073741824xf32> -> f32 + %rms0 = scalar.mulf %raw0, %rms_scale : f32 + %rms1 = scalar.mulf %raw1, %rms_scale : f32 + %w0 = view.load %weight_view[%off_a] : view<65536xf32> -> f32 + %w1 = view.load %weight_view[%off_b] : view<65536xf32> -> f32 + %x0 = scalar.mulf %rms0, %w0 : f32 + %x1 = scalar.mulf %rms1, %w1 : f32 + + %rot0, %rot1 = template.apply<@ggml.rope_f32.rotate_scalar_pair>(%x0, %x1, %cos_theta, %sin_theta) : (f32, f32, f32, f32) -> (f32, f32) + %out0 = scf.select %rotate, %rot0, %x0 : f32 + %out1 = scf.select %rotate, %rot1, %x1 : f32 + + %db1 = index.mul %i1, %d1 : index + %db2 = index.mul %i2, %d2 : index + %db3 = index.mul %i3, %d3 : index + %db01 = index.add %db1, %db2 : index + %dst_row = index.add %db01, %db3 : index + %da0 = index.add %dst_row, %off_a : index + %da = index.assume %da0 [range(%da0, 0, 1073741823)] : index + %db0 = index.add %dst_row, %off_b : index + %db = index.assume %db0 [range(%db0, 0, 1073741823)] : index + view.store %out0, %dst_view[%da] : f32, view<1073741824xf32> + view.store %out1, %dst_view[%db] : f32, view<1073741824xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom new file mode 100644 index 000000000000..ace56e1c8440 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_binary.loom @@ -0,0 +1,400 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.apply(%arg0: f32, %arg1: f32) -> (f32) + +template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) + +template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) + +template.decl @ggml.normalization.publish_f16.k16_vector4(%values: vector<4xf32>, %channel: index, %token: index, %token_count: index, %hidden_size: index, %packet: index, %output: buffer) + +template.decl @ggml.normalization.publish_q8.binary_vector4(%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 @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +config.decl @ggml.rmsnorm_binary_q8_1_x4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_q8_1_x4.op : %value: index where [range(%value, 0, 3)] + +template.decl @ggml.rmsnorm_binary_f32.body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +template.decl @ggml.rmsnorm_binary_f32.lowtoken_striped_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +template.decl @ggml.rmsnorm_binary_strided_f32.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +template.def<@ggml.rmsnorm_binary_f32.body> device @ggml_rmsnorm_binary_f32_body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %input_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %f16_output_view = buffer.view %f16_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf16> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %rhs_value = view.load %rhs_view[%channel] : view<[%hidden_size]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %result = template.apply<@ggml.binary_f32.apply>(%op, %normalized, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %publish_f16 { + %half = scalar.fptrunc %result : f32 to f16 + view.store %half, %f16_output_view[%token, %channel] : f16, view<[%launch_token_count]x[%hidden_size]xf16> + } + } + } + template.return +} + +template.def<@ggml.rmsnorm_binary_f32.lowtoken_striped_body> device @ggml_rmsnorm_binary_f32_lowtoken_striped_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %group = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %c0_offset = index.constant 0 : offset + %rounded_hidden = index.add %hidden_size, %c1023 : index + %stripe_count = index.div %rounded_hidden, %c1024 : index + %token0 = index.div %group, %stripe_count : index + %stripe = index.rem %group, %stripe_count : index + %stripe_base = index.mul %stripe, %c1024 : index + %stripe_limit0 = index.add %stripe_base, %c1024 : index + %stripe_limit = index.min %stripe_limit0, %hidden_size : index + %input_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %valid_token { + %first_channel = index.add %stripe_base, %workitem : index + scf.for %channel = [%first_channel to %stripe_limit step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %rhs_value = view.load %rhs_view[%channel] : view<[%hidden_size]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %result = template.apply<@ggml.binary_f32.apply>(%op, %normalized, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + template.return +} + +config.decl @ggml.rmsnorm_binary_f32.input_stride : %value: index where [range(%value, 128, 1048576)] + +template.def<@ggml.rmsnorm_binary_strided_f32.body> device @ggml_rmsnorm_binary_strided_f32_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %input_stride0 = config.get @ggml.rmsnorm_binary_f32.input_stride : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %hidden_size1 = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %input_stride1 = index.assume %input_stride0 [range(%input_stride0, 128, 1048576)] : index + %hidden_size, %input_stride = index.assume %hidden_size1, %input_stride1 [le(%hidden_size1, %input_stride1)] : index, index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0_offset = index.constant 0 : offset + %c256 = index.constant 256 : index + %input_noalias, %rhs_noalias, %output_noalias = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %input_stride, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%input_stride]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden_size step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%input_stride]xf32> -> f32 + %rhs_value = view.load %rhs_view[%channel] : view<[%hidden_size]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %result = template.apply<@ggml.binary_f32.apply>(%op, %normalized, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden_size]xf32> + } + } + template.return +} + +template.decl @ggml.rmsnorm_binary_f32_k16.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +template.def<@ggml.rmsnorm_binary_f32_k16.body> device @ggml_rmsnorm_binary_f32_k16_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) { + %hidden0 = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 1024, 8192), mul(%hidden0, 1024)] : index + %epsilon = config.get @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_f32.op : index + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %token_count)] : index + %lane = kernel.subgroup.lane.id : index + %wave = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %zero = scalar.constant 0.0 : f32 + %cache_zero = vector.constant 0.0 : vector<32xf32> + %input_na, %rhs_na, %output_na = buffer.assume.noalias %input, %rhs, %output : buffer, buffer, buffer + %input_view = buffer.view %input_na[%base] : buffer -> view<[%token_count]x[%hidden]xf32> + %rhs_view = buffer.view %rhs_na[%base] : buffer -> view<[%hidden]xf32> + %output_view = buffer.view %output_na[%base] : buffer -> view<[%token_count]x[%hidden]xf32> + %first_channel = index.mul %tid, %c2 : index + %stripe_count = index.div %hidden, %c256 : index + %sum0, %sum1, %cache0, %cache1, %cache2, %cache3 = scf.for %stripe = [%c0 to %stripe_count step %c1](%running0 = %zero : f32, %running1 = %zero : f32, %values0 = %cache_zero : vector<32xf32>, %values1 = %cache_zero : vector<32xf32>, %values2 = %cache_zero : vector<32xf32>, %values3 = %cache_zero : vector<32xf32>) -> (f32, f32, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) unroll { + %channel0 = index.madd %stripe, %c256, %first_channel : index + %channel = index.assume %channel0 [range(%channel0, 0, 8190), mul(%channel0, 2), lt(%channel0, %hidden)] : index + %input2 = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden]xf32> -> vector<2xf32> + %cache_phase = index.div %stripe, %c16 : index + %cache_stripe = index.rem %stripe, %c16 : index + %cache_base = index.mul %cache_stripe, %c2 : index + %value0 = vector.extract %input2[%c0] : vector<2xf32> -> f32 + %square0 = scalar.mulf %value0, %value0 : f32 + %next_sum0 = scalar.addf %running0, %square0 : f32 + %cache_index0 = index.add %cache_base, %c0 : index + %value1 = vector.extract %input2[%c1] : vector<2xf32> -> f32 + %square1 = scalar.mulf %value1, %value1 : f32 + %next_sum1 = scalar.addf %running1, %square1 : f32 + %cache_index1 = index.add %cache_base, %c1 : index + %phase0 = index.cmp eq, %cache_phase, %c0 : index + %next_values0 = scf.if %phase0 -> (vector<32xf32>) { + %insert0_0 = vector.insert %value0 into %values0[%cache_index0] : f32, vector<32xf32> + %insert0_1 = vector.insert %value1 into %insert0_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert0_1 : vector<32xf32> + } else { + scf.yield %values0 : vector<32xf32> + } + %phase1 = index.cmp eq, %cache_phase, %c1 : index + %next_values1 = scf.if %phase1 -> (vector<32xf32>) { + %insert1_0 = vector.insert %value0 into %values1[%cache_index0] : f32, vector<32xf32> + %insert1_1 = vector.insert %value1 into %insert1_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert1_1 : vector<32xf32> + } else { + scf.yield %values1 : vector<32xf32> + } + %phase2 = index.cmp eq, %cache_phase, %c2 : index + %next_values2 = scf.if %phase2 -> (vector<32xf32>) { + %insert2_0 = vector.insert %value0 into %values2[%cache_index0] : f32, vector<32xf32> + %insert2_1 = vector.insert %value1 into %insert2_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert2_1 : vector<32xf32> + } else { + scf.yield %values2 : vector<32xf32> + } + %phase3 = index.cmp eq, %cache_phase, %c3 : index + %next_values3 = scf.if %phase3 -> (vector<32xf32>) { + %insert3_0 = vector.insert %value0 into %values3[%cache_index0] : f32, vector<32xf32> + %insert3_1 = vector.insert %value1 into %insert3_0[%cache_index1] : f32, vector<32xf32> + scf.yield %insert3_1 : vector<32xf32> + } else { + scf.yield %values3 : vector<32xf32> + } + %fence_pos = index.rem %stripe, %c4 : index + %end_packet = index.cmp eq, %fence_pos, %c3 : index + scf.if %end_packet { + scf.schedule.fence + } + scf.yield %next_sum0, %next_sum1, %next_values0, %next_values1, %next_values2, %next_values3 : f32, f32, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32> + } + // Two original lane sums form a pair; sixteen threads form one original wave. + %pair_sum = scalar.addf %sum0, %sum1 : f32 + %shuffle_width = scalar.constant 32 : i32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %peer1, %valid1 = kernel.subgroup.shuffle %pair_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum4 = scalar.addf %pair_sum, %peer1 : f32 + %peer2, %valid2 = kernel.subgroup.shuffle %sum4, %xor2, %shuffle_width : f32, i32, i32 + %sum8 = scalar.addf %sum4, %peer2 : f32 + %peer4, %valid4 = kernel.subgroup.shuffle %sum8, %xor4, %shuffle_width : f32, i32, i32 + %sum16 = scalar.addf %sum8, %peer4 : f32 + %peer8, %valid8 = kernel.subgroup.shuffle %sum16, %xor8, %shuffle_width : f32, i32, i32 + %original_wave_sum = scalar.addf %sum16, %peer8 : f32 + %original_wave = index.div %tid, %c16 : index + %cohort_lane = index.rem %tid, %c16 : index + %cohort_leader = index.cmp eq, %cohort_lane, %c0 : index + %scratch_bytes = index.constant 1024 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%base] : buffer -> view<256xf32> + scf.if %cohort_leader { + view.store %original_wave_sum, %scratch_view[%original_wave] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %first_wave = index.cmp eq, %wave, %c0 : index + %live_partial = index.cmp ult, %lane, %c8 : index + %loads_partial = scalar.andi %first_wave, %live_partial : i1 + %partial = scf.if %loads_partial -> (f32) { + %loaded = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %loaded : f32 + } else { + scf.yield %zero : f32 + } + %row_sum = kernel.subgroup.reduce %partial : f32 + %leader = index.cmp eq, %tid, %c0 : index + scf.if %leader { + %hidden_i32 = index.cast %hidden : index to i32 + %hidden_f32 = scalar.sitofp %hidden_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + %pair_count = index.div %stripe_count, %c2 : index + %parity = index.rem %tid, %c2 : index + %odd = index.cmp ne, %parity, %c0 : index + %aligned_tid = index.sub %tid, %parity : index + %aligned_channel = index.mul %aligned_tid, %c2 : index + %stripe_lane = index.mul %parity, %c256 : index + %output_origin = index.add %aligned_channel, %stripe_lane : index + %c512 = index.constant 512 : index + scf.for %pair = [%c0 to %pair_count step %c1] unroll { + %channel0 = index.madd %pair, %c512, %output_origin : index + %channel = index.assume %channel0 [range(%channel0, 0, 8188), mul(%channel0, 4), lt(%channel0, %hidden)] : index + %stripe = index.mul %pair, %c2 : index + %phase = index.div %stripe, %c16 : index + %phase0 = index.cmp eq, %phase, %c0 : index + %phase1 = index.cmp eq, %phase, %c1 : index + %phase2 = index.cmp eq, %phase, %c2 : index + %values23 = scf.select %phase2, %cache2, %cache3 : vector<32xf32> + %values123 = scf.select %phase1, %cache1, %values23 : vector<32xf32> + %values = scf.select %phase0, %cache0, %values123 : vector<32xf32> + %cache_stripe = index.rem %stripe, %c16 : index + %cache_base = index.mul %cache_stripe, %c2 : index + %cache_index0 = index.add %cache_base, %c0 : index + %cached0 = vector.extract %values[%cache_index0] : vector<32xf32> -> f32 + %cache_index1 = index.add %cache_base, %c1 : index + %cached1 = vector.extract %values[%cache_index1] : vector<32xf32> -> f32 + %cache_index2 = index.add %cache_base, %c2 : index + %cached2 = vector.extract %values[%cache_index2] : vector<32xf32> -> f32 + %cache_index3 = index.add %cache_base, %c3 : index + %cached3 = vector.extract %values[%cache_index3] : vector<32xf32> -> f32 + %send0 = scf.select %odd, %cached0, %cached2 : f32 + %send1 = scf.select %odd, %cached1, %cached3 : f32 + %received0, %receive_valid0 = kernel.subgroup.shuffle %send0, %xor1, %shuffle_width : f32, i32, i32 + %received1, %receive_valid1 = kernel.subgroup.shuffle %send1, %xor1, %shuffle_width : f32, i32, i32 + %value0 = scf.select %odd, %received0, %cached0 : f32 + %value1 = scf.select %odd, %received1, %cached1 : f32 + %value2 = scf.select %odd, %cached2, %received0 : f32 + %value3 = scf.select %odd, %cached3, %received1 : f32 + %weights = vector.load %rhs_view[%channel] : view<[%hidden]xf32> -> vector<4xf32> + %result_zero = vector.constant 0.0 : vector<4xf32> + %weight0 = vector.extract %weights[%c0] : vector<4xf32> -> f32 + %norm0 = template.apply<@ggml.rmsnorm_f32.apply>(%value0, %row_scale) : (f32, f32) -> (f32) + %out0 = template.apply<@ggml.binary_f32.apply>(%op, %norm0, %weight0) : (index, f32, f32) -> (f32) + %result0 = vector.insert %out0 into %result_zero[%c0] : f32, vector<4xf32> + %weight1 = vector.extract %weights[%c1] : vector<4xf32> -> f32 + %norm1 = template.apply<@ggml.rmsnorm_f32.apply>(%value1, %row_scale) : (f32, f32) -> (f32) + %out1 = template.apply<@ggml.binary_f32.apply>(%op, %norm1, %weight1) : (index, f32, f32) -> (f32) + %result1 = vector.insert %out1 into %result0[%c1] : f32, vector<4xf32> + %weight2 = vector.extract %weights[%c2] : vector<4xf32> -> f32 + %norm2 = template.apply<@ggml.rmsnorm_f32.apply>(%value2, %row_scale) : (f32, f32) -> (f32) + %out2 = template.apply<@ggml.binary_f32.apply>(%op, %norm2, %weight2) : (index, f32, f32) -> (f32) + %result2 = vector.insert %out2 into %result1[%c2] : f32, vector<4xf32> + %weight3 = vector.extract %weights[%c3] : vector<4xf32> -> f32 + %norm3 = template.apply<@ggml.rmsnorm_f32.apply>(%value3, %row_scale) : (f32, f32) -> (f32) + %out3 = template.apply<@ggml.binary_f32.apply>(%op, %norm3, %weight3) : (index, f32, f32) -> (f32) + %result3 = vector.insert %out3 into %result2[%c3] : f32, vector<4xf32> + vector.store %result3, %output_view[%token, %channel] : vector<4xf32>, view<[%token_count]x[%hidden]xf32> + template.apply<@ggml.normalization.publish_f16.k16_vector4>(%result3, %channel, %token, %token_count, %hidden, %pair, %f16_output) : (vector<4xf32>, index, index, index, index, index, buffer) + } + template.return +} + +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) + +template.def<@ggml.rmsnorm_binary_q8_1_x4.body> device @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) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_q8_1_x4.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_q8_1_x4.rms_epsilon : f32 + %op = config.get @ggml.rmsnorm_binary_q8_1_x4.op : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %dispatch_group = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %tile_round = index.constant 1023 : index + %decode_limit = index.constant 6 : index + %rounded_hidden = index.add %hidden_size, %tile_round : index + %stripe_count = index.div %rounded_hidden, %c1024 : index + %is_decode = index.cmp ult, %token_count, %decode_limit : index + %groups_per_row = scf.select %is_decode, %stripe_count, %c1 : index + %token0 = index.div %dispatch_group, %groups_per_row : index + %stripe_id = index.rem %dispatch_group, %groups_per_row : index + %stripe_start = index.mul %stripe_id, %c1024 : index + %stripe_step = index.mul %groups_per_row, %c1024 : index + %group_bytes = index.constant 144 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %input_noalias, %rhs_noalias, %output_noalias, %q8_output_noalias = buffer.assume.noalias %input, %rhs, %output, %q8_output : buffer, buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %f16_output_view = buffer.view %f16_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf16> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + scf.for %stripe_base = [%stripe_start to %hidden_size step %stripe_step] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word0 = index.cmp ult, %channel, %hidden_size : index + %valid_word = scalar.andi %valid_word0, %valid_token : i1 + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %rhs_values = vector.load.mask %rhs_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %result0 = vector.extract %normalized[0] : vector<4xf32> -> f32 + %result1 = vector.extract %normalized[1] : vector<4xf32> -> f32 + %result2 = vector.extract %normalized[2] : vector<4xf32> -> f32 + %result3 = vector.extract %normalized[3] : vector<4xf32> -> f32 + %rhs0 = vector.extract %rhs_values[0] : vector<4xf32> -> f32 + %rhs1 = vector.extract %rhs_values[1] : vector<4xf32> -> f32 + %rhs2 = vector.extract %rhs_values[2] : vector<4xf32> -> f32 + %rhs3 = vector.extract %rhs_values[3] : vector<4xf32> -> f32 + %binary0 = template.apply<@ggml.binary_f32.apply>(%op, %result0, %rhs0) : (index, f32, f32) -> (f32) + %binary1 = template.apply<@ggml.binary_f32.apply>(%op, %result1, %rhs1) : (index, f32, f32) -> (f32) + %binary2 = template.apply<@ggml.binary_f32.apply>(%op, %result2, %rhs2) : (index, f32, f32) -> (f32) + %binary3 = template.apply<@ggml.binary_f32.apply>(%op, %result3, %rhs3) : (index, f32, f32) -> (f32) + %packed0 = vector.insert %binary0 into %c0_f32x4[0] : f32, vector<4xf32> + %packed1 = vector.insert %binary1 into %packed0[1] : f32, vector<4xf32> + %packed2 = vector.insert %binary2 into %packed1[2] : f32, vector<4xf32> + %packed = vector.insert %binary3 into %packed2[3] : f32, vector<4xf32> + vector.store.mask %packed, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + scf.if %publish_f16 { + %half = vector.fptrunc %packed : vector<4xf32> to vector<4xf16> + vector.store.mask %half, %f16_output_view[%token, %channel], %mask : vector<4xf16>, view<[%launch_token_count]x[%hidden_size]xf16>, vector<4xi1> + } + template.apply<@ggml.normalization.publish_q8.binary_vector4>(%valid_word, %token_output_byte_base, %channel, %packed, %scratch_values, %scratch_d, %q8_output_noalias) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom new file mode 100644 index 000000000000..d80f1546660f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_gate.loom @@ -0,0 +1,62 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.subgroup_row_scale(%token_count: index, %token: index, %hidden_size: index, %epsilon: f32, %input: buffer) -> (f32) +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) +template.decl @ggml.normalization.publish_f16.gate_vector4(%values: vector<4xf32>, %row: index, %channel: index, %token_count: index, %hidden_size: index, %row_width: index, %output: buffer) +template.decl @ggml.normalization.publish_q8.gate_vector4(%publish_word: i1, %row: index, %hidden_size: index, %channel: index, %values: vector<4xf32>, %output: buffer) +template.decl @ggml.rmsnorm_gate_f32.publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) + +config.decl @ggml.rmsnorm_gate_f32.hidden_size : %value: index where [range(%value, 128, 1024), mul(%value, 128)] +config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] +config.decl @ggml.rmsnorm_gate_f32.f16_output_row_width : %value: index where [range(%value, 0, 32768), mul(%value, 128)] + +template.def<@ggml.rmsnorm_gate_f32.publish_body> device @ggml_rmsnorm_gate_f32_publish_body(%publish_q8: i1, %token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %alternate_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1048576)] : index + %hidden0 = config.get @ggml.rmsnorm_gate_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 1024), mul(%hidden0, 128)] : index + %epsilon = config.get @ggml.rmsnorm_gate_f32.rms_epsilon : f32 + %gate_op = config.get @ggml.rmsnorm_gate_f32.gate_op : index + %row_width = config.get @ggml.rmsnorm_gate_f32.f16_output_row_width : index + %z = index.constant 0 : index + %zero_offset = index.constant 0 : offset + %four = index.constant 4 : index + %eight = index.constant 8 : index + %c128 = index.constant 128 : index + %wg = kernel.workgroup.id : index + %wave = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %row0 = index.madd %wg, %eight, %wave : index + %valid = index.cmp ult, %row0, %bounded_token_count : index + %safe = scf.select %valid, %row0, %z : index + %row = index.assume %safe [lt(%safe, %bounded_token_count)] : index + %channel0 = index.mul %lane, %four : index + %input_n, %weight_n, %raw_gate_n, %output_n, %alternate_n = buffer.assume.noalias %input, %weight, %raw_gate, %output, %alternate_output : buffer, buffer, buffer, buffer, buffer + %iv = buffer.view %input_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %wv = buffer.view %weight_n[%zero_offset] : buffer -> view<[%hidden]xf32> + %gv = buffer.view %raw_gate_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %ov = buffer.view %output_n[%zero_offset] : buffer -> view<[%bounded_token_count]x[%hidden]xf32> + %scale = template.apply<@ggml.rmsnorm_f32.subgroup_row_scale>(%bounded_token_count, %row, %hidden, %epsilon, %input_n) : (index, index, index, f32, buffer) -> (f32) + %scale_v = vector.splat %scale : vector<4xf32> + %last_channel = index.sub %hidden, %four : index + scf.if %valid { + scf.for %channel = [%channel0 to %hidden step %c128] { + %bounded_channel = index.assume %channel [range(%channel, 0, 1020), mul(%channel, 4), le(%channel, %last_channel)] : index + %x = vector.load %iv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> + %w = vector.load %wv[%bounded_channel] : view<[%hidden]xf32> -> vector<4xf32> + %g = vector.load %gv[%row, %bounded_channel] : view<[%bounded_token_count]x[%hidden]xf32> -> vector<4xf32> + %norm = vector.mulf %x, %scale_v : vector<4xf32> + %side = vector.mulf %norm, %w : vector<4xf32> + %activated = template.apply<@ggml.unary_f32.apply_vector4>(%gate_op, %g) : (index, vector<4xf32>) -> (vector<4xf32>) + %out = vector.mulf %side, %activated : vector<4xf32> + vector.store %out, %ov[%row, %bounded_channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden]xf32> + scf.if %publish_q8 { + template.apply<@ggml.normalization.publish_q8.gate_vector4>(%valid, %row, %hidden, %bounded_channel, %out, %alternate_n) : (i1, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.normalization.publish_f16.gate_vector4>(%out, %row, %bounded_channel, %bounded_token_count, %hidden, %row_width, %alternate_n) : (vector<4xf32>, index, index, index, index, index, buffer) + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom new file mode 100644 index 000000000000..c3b9c47a181d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom @@ -0,0 +1,182 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_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 + +// Shared RMSNorm and GGML Q8_1 x4 row producer. The caller owns launch geometry +// and supplies the logical token whose complete hidden row this workgroup owns. +// The attention export discards the ordinary F32 row, while feed-forward +// publishes it for the router without rereading and requantizing the normalized +// values. +template.def<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body> device @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %physical_group_count = index.div %hidden_size, %c128 : index + %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + // Reduce the complete row before any block-local quantization. + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %value, %value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %row_scale_vector = vector.splat %row_scale : vector<4xf32> + %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 + // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need + // only be divisible by 128; inactive workitems in the final stripe carry + // zeros and never publish. + scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { + %word_element_add = index.mul %workitem, %c4 : index + %channel = index.add %stripe_base, %word_element_add : index + %valid_word = index.cmp ult, %channel, %hidden_size : index + %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> + %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> + %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> + %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> + scf.if %publishes_normalized { + vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> + } + %absolute_values = vector.absf %normalized : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %word_in_block = index.rem %workitem, %c8 : index + %block_in_stripe = index.div %workitem, %c8 : index + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 + scf.if %writes_block_d { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = scf.if %valid_word -> (f32) { + %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 + scf.yield %block_d : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 + scf.if %publishes_q8_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 + scf.if %publishes_block_ds { + %cohort_base = index.mul %block_in_stripe, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_symmetric_i4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_symmetric_i4.loom new file mode 100644 index 000000000000..ffcd900874a4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/normalization/rmsnorm_symmetric_i4.loom @@ -0,0 +1,496 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.row_scale(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5: buffer) -> (f32, index, index, i1) + +config.decl @ggml.add_rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon : f32 +config.decl @ggml.rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_symmetric_i4.rms_epsilon : f32 +config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon : f32 + +template.decl @ggml.rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.def<@ggml.rmsnorm_binary_symmetric_i4_k32.body> device @ggml_rmsnorm_binary_symmetric_i4_k32_body(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_binary_symmetric_i4.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_binary_symmetric_i4.rms_epsilon : f32 + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %scratch_bytes = index.constant 1024 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %amax_epsilon = scalar.constant 1.0000000000000001e-30 : f32 + %one_seventh = scalar.constant 0.14285714285714285 : f32 + %seven = scalar.constant 7.0 : f32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %shuffle_width = scalar.constant 32 : i32 + %shift4 = scalar.constant 4 : i32 + %shift8 = scalar.constant 8 : i32 + %shift12 = scalar.constant 12 : i32 + %nibble_mask = vector.constant 15 : vector<4xi32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token = index.assume %safe_token0 [lt(%safe_token0, %bounded_token_count)] : index + + %input_noalias, %weight_noalias, %output_noalias, %i4_qs_noalias, %i4_ds_noalias, %i4_sums_noalias = buffer.assume.noalias %input, %weight, %output, %i4_qs, %i4_ds, %i4_sums : buffer, buffer, buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %element_count = index.mul %bounded_token_count, %hidden_size : index + %group32_count = index.div %element_count, %c32 : index + %group64_count_per_row = index.div %hidden_size, %c64 : index + %qs_halfword_count = index.div %element_count, %c4 : index + %i4_qs_view = buffer.view %i4_qs_noalias[%c0_offset] : buffer -> view<[%qs_halfword_count]xi16> + %i4_ds_view = buffer.view %i4_ds_noalias[%c0_offset] : buffer -> view<[%group32_count]xf32> + %i4_sums_view = buffer.view %i4_sums_noalias[%c0_offset] : buffer -> view<[%group32_count]xi32> + + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %input_value = view.load %input_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> f32 + %square = scalar.mulf %input_value, %input_value : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %row_scale = scalar.rsqrtf %biased_mean : f32 + view.store %row_scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + + %wave = index.div %workitem, %c32 : index + %group_in_wave = index.div %lane, %c16 : index + %lane_in_group = index.rem %lane, %c16 : index + %group_half = index.div %lane_in_group, %c8 : index + %lane_in_half = index.rem %lane_in_group, %c8 : index + %is_scale_leader = index.cmp eq, %lane_in_group, %c0 : index + %is_sum_leader = index.cmp eq, %lane_in_half, %c0 : index + %wave_group_base = index.mul %wave, %c2 : index + %local_group = index.add %wave_group_base, %group_in_wave : index + %token_group_base = index.mul %token, %group64_count_per_row : index + scf.if %valid_token { + scf.for %group_base = [%c0 to %group64_count_per_row step %c16] { + %group_in_row = index.add %group_base, %local_group : index + %in_range = index.cmp ult, %group_in_row, %group64_count_per_row : index + scf.if %in_range { + %channel_base = index.mul %group_in_row, %c64 : index + %lane_channel = index.mul %lane_in_group, %c4 : index + %channel = index.add %channel_base, %lane_channel : index + %input_values = vector.load %input_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %learned_weights = vector.load %weight_view[%channel] : view<[%hidden_size]xf32> -> vector<4xf32> + %scale_vector = vector.splat %row_scale : vector<4xf32> + %normalized = vector.mulf %input_values, %scale_vector : vector<4xf32> + %result = vector.mulf %normalized, %learned_weights : vector<4xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden_size]xf32> + + %absolute_values = vector.absf %result : vector<4xf32> + %lane_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %max_x1_peer, %max_x1_valid = kernel.subgroup.shuffle %lane_max, %xor1, %shuffle_width : f32, i32, i32 + %max_x1 = scalar.maxnumf %lane_max, %max_x1_peer : f32 + %max_x2_peer, %max_x2_valid = kernel.subgroup.shuffle %max_x1, %xor2, %shuffle_width : f32, i32, i32 + %max_x2 = scalar.maxnumf %max_x1, %max_x2_peer : f32 + %max_x4_peer, %max_x4_valid = kernel.subgroup.shuffle %max_x2, %xor4, %shuffle_width : f32, i32, i32 + %max_x4 = scalar.maxnumf %max_x2, %max_x4_peer : f32 + %max_x8_peer, %max_x8_valid = kernel.subgroup.shuffle %max_x4, %xor8, %shuffle_width : f32, i32, i32 + %group_max = scalar.maxnumf %max_x4, %max_x8_peer : f32 + %amax = scalar.maxnumf %group_max, %amax_epsilon : f32 + %activation_scale = scalar.mulf %amax, %one_seventh : f32 + %activation_rscale = scalar.divf %seven, %amax : f32 + %group64 = index.add %token_group_base, %group_in_row : index + %group32_base = index.mul %group64, %c2 : index + %group32_high = index.add %group32_base, %c1 : index + scf.if %is_scale_leader { + view.store %activation_scale, %i4_ds_view[%group32_base] : f32, view<[%group32_count]xf32> + view.store %activation_scale, %i4_ds_view[%group32_high] : f32, view<[%group32_count]xf32> + } + + %rscale_vector = vector.splat %activation_rscale : vector<4xf32> + %scaled = vector.mulf %result, %rscale_vector : vector<4xf32> + %rounded = vector.roundf %scaled : vector<4xf32> + %quantized_i32 = vector.fptosi %rounded : vector<4xf32> to vector<4xi32> + %quantized_nibbles = vector.andi %quantized_i32, %nibble_mask : vector<4xi32> + %q0 = vector.extract %quantized_nibbles[0] : vector<4xi32> -> i32 + %q1 = vector.extract %quantized_nibbles[1] : vector<4xi32> -> i32 + %q2 = vector.extract %quantized_nibbles[2] : vector<4xi32> -> i32 + %q3 = vector.extract %quantized_nibbles[3] : vector<4xi32> -> i32 + %q1_shifted = scalar.shli %q1, %shift4 : i32 + %q2_shifted = scalar.shli %q2, %shift8 : i32 + %q3_shifted = scalar.shli %q3, %shift12 : i32 + %q01 = scalar.ori %q0, %q1_shifted : i32 + %q23 = scalar.ori %q2_shifted, %q3_shifted : i32 + %packed_i32 = scalar.ori %q01, %q23 : i32 + %packed_i16 = scalar.trunci %packed_i32 : i32 to i16 + %word_base = index.mul %group64, %c16 : index + %word_index = index.add %word_base, %lane_in_group : index + view.store %packed_i16, %i4_qs_view[%word_index] : i16, view<[%qs_halfword_count]xi16> + + %lane_sum = vector.reduce %rounded, %c0_f32 : vector<4xf32>, f32 + %sum_x1_peer, %sum_x1_valid = kernel.subgroup.shuffle %lane_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum_x1 = scalar.addf %lane_sum, %sum_x1_peer : f32 + %sum_x2_peer, %sum_x2_valid = kernel.subgroup.shuffle %sum_x1, %xor2, %shuffle_width : f32, i32, i32 + %sum_x2 = scalar.addf %sum_x1, %sum_x2_peer : f32 + %sum_x4_peer, %sum_x4_valid = kernel.subgroup.shuffle %sum_x2, %xor4, %shuffle_width : f32, i32, i32 + %group_sum = scalar.addf %sum_x2, %sum_x4_peer : f32 + scf.if %is_sum_leader { + %sum_value = scalar.fptosi %group_sum : f32 to i32 + %sum_index = index.add %group32_base, %group_half : index + view.store %sum_value, %i4_sums_view[%sum_index] : i32, view<[%group32_count]xi32> + } + } + } + } + template.return +} + + +template.decl @ggml.add_rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.def<@ggml.add_rmsnorm_binary_symmetric_i4_k32.body> device @ggml_add_rmsnorm_binary_symmetric_i4_k32_body(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { + %hidden_size0 = config.get @ggml.add_rmsnorm_binary_symmetric_i4.hidden_size : index + %epsilon = config.get @ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon : f32 + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %scratch_bytes = index.constant 1024 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %amax_epsilon = scalar.constant 1.0000000000000001e-30 : f32 + %one_seventh = scalar.constant 0.14285714285714285 : f32 + %seven = scalar.constant 7.0 : f32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %shuffle_width = scalar.constant 32 : i32 + %shift4 = scalar.constant 4 : i32 + %shift8 = scalar.constant 8 : i32 + %shift12 = scalar.constant 12 : i32 + %nibble_mask = vector.constant 15 : vector<4xi32> + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %token = index.assume %safe_token0 [lt(%safe_token0, %bounded_token_count)] : index + + %lhs_noalias, %rhs_noalias, %residual_noalias, %weight_noalias, %output_noalias, %i4_qs_noalias, %i4_ds_noalias, %i4_sums_noalias = buffer.assume.noalias %lhs, %rhs, %residual_output, %weight, %output, %i4_qs, %i4_ds, %i4_sums : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer + %lhs_view = buffer.view %lhs_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %rhs_view = buffer.view %rhs_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %residual_view = buffer.view %residual_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %element_count = index.mul %bounded_token_count, %hidden_size : index + %group32_count = index.div %element_count, %c32 : index + %group64_count_per_row = index.div %hidden_size, %c64 : index + %qs_halfword_count = index.div %element_count, %c4 : index + %i4_qs_view = buffer.view %i4_qs_noalias[%c0_offset] : buffer -> view<[%qs_halfword_count]xi16> + %i4_ds_view = buffer.view %i4_ds_noalias[%c0_offset] : buffer -> view<[%group32_count]xf32> + %i4_sums_view = buffer.view %i4_sums_noalias[%c0_offset] : buffer -> view<[%group32_count]xi32> + + %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %lhs_value = view.load %lhs_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> f32 + %rhs_value = view.load %rhs_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> f32 + %residual = scalar.addf %lhs_value, %rhs_value : f32 + view.store %residual, %residual_view[%token, %channel] : f32, view<[%bounded_token_count]x[%hidden_size]xf32> + %square = scalar.mulf %residual, %residual : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %row_scale = scalar.rsqrtf %biased_mean : f32 + view.store %row_scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + kernel.barrier scope(workgroup) ordering(acq_rel) + + %wave = index.div %workitem, %c32 : index + %group_in_wave = index.div %lane, %c16 : index + %lane_in_group = index.rem %lane, %c16 : index + %group_half = index.div %lane_in_group, %c8 : index + %lane_in_half = index.rem %lane_in_group, %c8 : index + %is_scale_leader = index.cmp eq, %lane_in_group, %c0 : index + %is_sum_leader = index.cmp eq, %lane_in_half, %c0 : index + %wave_group_base = index.mul %wave, %c2 : index + %local_group = index.add %wave_group_base, %group_in_wave : index + %token_group_base = index.mul %token, %group64_count_per_row : index + scf.if %valid_token { + scf.for %group_base = [%c0 to %group64_count_per_row step %c16] { + %group_in_row = index.add %group_base, %local_group : index + %in_range = index.cmp ult, %group_in_row, %group64_count_per_row : index + scf.if %in_range { + %channel_base = index.mul %group_in_row, %c64 : index + %lane_channel = index.mul %lane_in_group, %c4 : index + %channel = index.add %channel_base, %lane_channel : index + %residual_values = vector.load %residual_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %learned_weights = vector.load %weight_view[%channel] : view<[%hidden_size]xf32> -> vector<4xf32> + %scale_vector = vector.splat %row_scale : vector<4xf32> + %normalized = vector.mulf %residual_values, %scale_vector : vector<4xf32> + %result = vector.mulf %normalized, %learned_weights : vector<4xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden_size]xf32> + + %absolute_values = vector.absf %result : vector<4xf32> + %lane_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %max_x1_peer, %max_x1_valid = kernel.subgroup.shuffle %lane_max, %xor1, %shuffle_width : f32, i32, i32 + %max_x1 = scalar.maxnumf %lane_max, %max_x1_peer : f32 + %max_x2_peer, %max_x2_valid = kernel.subgroup.shuffle %max_x1, %xor2, %shuffle_width : f32, i32, i32 + %max_x2 = scalar.maxnumf %max_x1, %max_x2_peer : f32 + %max_x4_peer, %max_x4_valid = kernel.subgroup.shuffle %max_x2, %xor4, %shuffle_width : f32, i32, i32 + %max_x4 = scalar.maxnumf %max_x2, %max_x4_peer : f32 + %max_x8_peer, %max_x8_valid = kernel.subgroup.shuffle %max_x4, %xor8, %shuffle_width : f32, i32, i32 + %group_max = scalar.maxnumf %max_x4, %max_x8_peer : f32 + %amax = scalar.maxnumf %group_max, %amax_epsilon : f32 + %activation_scale = scalar.mulf %amax, %one_seventh : f32 + %activation_rscale = scalar.divf %seven, %amax : f32 + %group64 = index.add %token_group_base, %group_in_row : index + %group32_base = index.mul %group64, %c2 : index + %group32_high = index.add %group32_base, %c1 : index + scf.if %is_scale_leader { + view.store %activation_scale, %i4_ds_view[%group32_base] : f32, view<[%group32_count]xf32> + view.store %activation_scale, %i4_ds_view[%group32_high] : f32, view<[%group32_count]xf32> + } + + %rscale_vector = vector.splat %activation_rscale : vector<4xf32> + %scaled = vector.mulf %result, %rscale_vector : vector<4xf32> + %rounded = vector.roundf %scaled : vector<4xf32> + %quantized_i32 = vector.fptosi %rounded : vector<4xf32> to vector<4xi32> + %quantized_nibbles = vector.andi %quantized_i32, %nibble_mask : vector<4xi32> + %q0 = vector.extract %quantized_nibbles[0] : vector<4xi32> -> i32 + %q1 = vector.extract %quantized_nibbles[1] : vector<4xi32> -> i32 + %q2 = vector.extract %quantized_nibbles[2] : vector<4xi32> -> i32 + %q3 = vector.extract %quantized_nibbles[3] : vector<4xi32> -> i32 + %q1_shifted = scalar.shli %q1, %shift4 : i32 + %q2_shifted = scalar.shli %q2, %shift8 : i32 + %q3_shifted = scalar.shli %q3, %shift12 : i32 + %q01 = scalar.ori %q0, %q1_shifted : i32 + %q23 = scalar.ori %q2_shifted, %q3_shifted : i32 + %packed_i32 = scalar.ori %q01, %q23 : i32 + %packed_i16 = scalar.trunci %packed_i32 : i32 to i16 + %word_base = index.mul %group64, %c16 : index + %word_index = index.add %word_base, %lane_in_group : index + view.store %packed_i16, %i4_qs_view[%word_index] : i16, view<[%qs_halfword_count]xi16> + + %lane_sum = vector.reduce %rounded, %c0_f32 : vector<4xf32>, f32 + %sum_x1_peer, %sum_x1_valid = kernel.subgroup.shuffle %lane_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum_x1 = scalar.addf %lane_sum, %sum_x1_peer : f32 + %sum_x2_peer, %sum_x2_valid = kernel.subgroup.shuffle %sum_x1, %xor2, %shuffle_width : f32, i32, i32 + %sum_x2 = scalar.addf %sum_x1, %sum_x2_peer : f32 + %sum_x4_peer, %sum_x4_valid = kernel.subgroup.shuffle %sum_x2, %xor4, %shuffle_width : f32, i32, i32 + %group_sum = scalar.addf %sum_x2, %sum_x4_peer : f32 + scf.if %is_sum_leader { + %sum_value = scalar.fptosi %group_sum : f32 to i32 + %sum_index = index.add %group32_base, %group_half : index + view.store %sum_value, %i4_sums_view[%sum_index] : i32, view<[%group32_count]xi32> + } + } + } + } + template.return +} + + +template.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +template.def<@ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body> device @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_body(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { + %hidden_size0 = config.get @ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size : index + %epsilon = config.get @ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon : f32 + %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %amax_epsilon = scalar.constant 1.0000000000000001e-30 : f32 + %one_seventh = scalar.constant 0.14285714285714285 : f32 + %seven = scalar.constant 7.0 : f32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %shuffle_width = scalar.constant 32 : i32 + %shift4 = scalar.constant 4 : i32 + %shift8 = scalar.constant 8 : i32 + %shift12 = scalar.constant 12 : i32 + %nibble_mask = vector.constant 15 : vector<4xi32> + + %input_noalias, %weight_noalias, %raw_gate_noalias, %output_noalias, %i4_qs_noalias, %i4_ds_noalias, %i4_sums_noalias = buffer.assume.noalias %input, %weight, %raw_gate, %output, %i4_qs, %i4_ds, %i4_sums : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %row_scale, %token, %bounded_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden_size, %hidden_size, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %raw_gate_view = buffer.view %raw_gate_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%hidden_size]xf32> + %element_count = index.mul %bounded_token_count, %hidden_size : index + %group32_count = index.div %element_count, %c32 : index + %group64_count_per_row = index.div %hidden_size, %c64 : index + %qs_halfword_count = index.div %element_count, %c4 : index + %i4_qs_view = buffer.view %i4_qs_noalias[%c0_offset] : buffer -> view<[%qs_halfword_count]xi16> + %i4_ds_view = buffer.view %i4_ds_noalias[%c0_offset] : buffer -> view<[%group32_count]xf32> + %i4_sums_view = buffer.view %i4_sums_noalias[%c0_offset] : buffer -> view<[%group32_count]xi32> + + %wave = index.div %workitem, %c32 : index + %group_in_wave = index.div %lane, %c16 : index + %lane_in_group = index.rem %lane, %c16 : index + %group_half = index.div %lane_in_group, %c8 : index + %lane_in_half = index.rem %lane_in_group, %c8 : index + %is_scale_leader = index.cmp eq, %lane_in_group, %c0 : index + %is_sum_leader = index.cmp eq, %lane_in_half, %c0 : index + %wave_group_base = index.mul %wave, %c2 : index + %local_group = index.add %wave_group_base, %group_in_wave : index + %token_group_base = index.mul %token, %group64_count_per_row : index + scf.if %valid_token { + scf.for %group_base = [%c0 to %group64_count_per_row step %c16] { + %group_in_row = index.add %group_base, %local_group : index + %in_range = index.cmp ult, %group_in_row, %group64_count_per_row : index + scf.if %in_range { + %channel_base = index.mul %group_in_row, %c64 : index + %lane_channel = index.mul %lane_in_group, %c4 : index + %channel = index.add %channel_base, %lane_channel : index + %values = vector.load %input_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %learned_weights = vector.load %weight_view[%channel] : view<[%hidden_size]xf32> -> vector<4xf32> + %raw_gate_values = vector.load %raw_gate_view[%token, %channel] : view<[%bounded_token_count]x[%hidden_size]xf32> -> vector<4xf32> + %scale_vector = vector.splat %row_scale : vector<4xf32> + %normalized = vector.mulf %values, %scale_vector : vector<4xf32> + %side = vector.mulf %normalized, %learned_weights : vector<4xf32> + %activated = vector.siluf %raw_gate_values : vector<4xf32> + %result = vector.mulf %side, %activated : vector<4xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%hidden_size]xf32> + + %absolute_values = vector.absf %result : vector<4xf32> + %lane_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %max_x1_peer, %max_x1_valid = kernel.subgroup.shuffle %lane_max, %xor1, %shuffle_width : f32, i32, i32 + %max_x1 = scalar.maxnumf %lane_max, %max_x1_peer : f32 + %max_x2_peer, %max_x2_valid = kernel.subgroup.shuffle %max_x1, %xor2, %shuffle_width : f32, i32, i32 + %max_x2 = scalar.maxnumf %max_x1, %max_x2_peer : f32 + %max_x4_peer, %max_x4_valid = kernel.subgroup.shuffle %max_x2, %xor4, %shuffle_width : f32, i32, i32 + %max_x4 = scalar.maxnumf %max_x2, %max_x4_peer : f32 + %max_x8_peer, %max_x8_valid = kernel.subgroup.shuffle %max_x4, %xor8, %shuffle_width : f32, i32, i32 + %group_max = scalar.maxnumf %max_x4, %max_x8_peer : f32 + %amax = scalar.maxnumf %group_max, %amax_epsilon : f32 + %activation_scale = scalar.mulf %amax, %one_seventh : f32 + %activation_rscale = scalar.divf %seven, %amax : f32 + %group64 = index.add %token_group_base, %group_in_row : index + %group32_base = index.mul %group64, %c2 : index + %group32_high = index.add %group32_base, %c1 : index + scf.if %is_scale_leader { + view.store %activation_scale, %i4_ds_view[%group32_base] : f32, view<[%group32_count]xf32> + view.store %activation_scale, %i4_ds_view[%group32_high] : f32, view<[%group32_count]xf32> + } + + %rscale_vector = vector.splat %activation_rscale : vector<4xf32> + %scaled = vector.mulf %result, %rscale_vector : vector<4xf32> + %rounded = vector.roundf %scaled : vector<4xf32> + %quantized_i32 = vector.fptosi %rounded : vector<4xf32> to vector<4xi32> + %quantized_nibbles = vector.andi %quantized_i32, %nibble_mask : vector<4xi32> + %q0 = vector.extract %quantized_nibbles[0] : vector<4xi32> -> i32 + %q1 = vector.extract %quantized_nibbles[1] : vector<4xi32> -> i32 + %q2 = vector.extract %quantized_nibbles[2] : vector<4xi32> -> i32 + %q3 = vector.extract %quantized_nibbles[3] : vector<4xi32> -> i32 + %q1_shifted = scalar.shli %q1, %shift4 : i32 + %q2_shifted = scalar.shli %q2, %shift8 : i32 + %q3_shifted = scalar.shli %q3, %shift12 : i32 + %q01 = scalar.ori %q0, %q1_shifted : i32 + %q23 = scalar.ori %q2_shifted, %q3_shifted : i32 + %packed_i32 = scalar.ori %q01, %q23 : i32 + %packed_i16 = scalar.trunci %packed_i32 : i32 to i16 + %word_base = index.mul %group64, %c16 : index + %word_index = index.add %word_base, %lane_in_group : index + view.store %packed_i16, %i4_qs_view[%word_index] : i16, view<[%qs_halfword_count]xi16> + + %lane_sum = vector.reduce %rounded, %c0_f32 : vector<4xf32>, f32 + %sum_x1_peer, %sum_x1_valid = kernel.subgroup.shuffle %lane_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum_x1 = scalar.addf %lane_sum, %sum_x1_peer : f32 + %sum_x2_peer, %sum_x2_valid = kernel.subgroup.shuffle %sum_x1, %xor2, %shuffle_width : f32, i32, i32 + %sum_x2 = scalar.addf %sum_x1, %sum_x2_peer : f32 + %sum_x4_peer, %sum_x4_valid = kernel.subgroup.shuffle %sum_x2, %xor4, %shuffle_width : f32, i32, i32 + %group_sum = scalar.addf %sum_x2, %sum_x4_peer : f32 + scf.if %is_sum_leader { + %sum_value = scalar.fptosi %group_sum : f32 to i32 + %sum_index = index.add %group32_base, %group_half : index + view.store %sum_value, %i4_sums_view[%sum_index] : i32, view<[%group32_count]xi32> + } + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/publish_f32.loom new file mode 100644 index 000000000000..006a83106659 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/publish_f32.loom @@ -0,0 +1,81 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.publish_f32.f16_vector4(%publish_word: i1, %token: index, %channel: index, %token_count: index, %hidden_size: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.publish_f32.f32_vector4(%publish_word: i1, %token: index, %channel: index, %token_count: index, %hidden_size: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.publish_f32.next_vector4(%format: index, %publish_word: i1, %token: index, %channel: index, %token_count: index, %hidden_size: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.publish_f32.q8_1_x4_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.quantize_q8_1_x4.publish_vector4(%arg0: i1, %arg1: offset, %arg2: index, %arg3: vector<4xf32>, %arg4: view<256xf32>, %arg5: view<32xf32>, %arg6: buffer) + +template.def<@ggml.publish_f32.f32_vector4> device @ggml_publish_f32_f32_vector4(%publish_word: i1, %token: index, %channel: index, %token_count: index, %hidden_size: index, %values: vector<4xf32>, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 128, 1073741824), mul(%hidden_size, 128)] : index + %c0_offset = index.constant 0 : offset + scf.if %publish_word { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + vector.store %values, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + } + template.return +} + +template.def<@ggml.publish_f32.f16_vector4> device @ggml_publish_f32_f16_vector4(%publish_word: i1, %token: index, %channel: index, %token_count: index, %hidden_size: index, %values: vector<4xf32>, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 128, 32768), mul(%hidden_size, 128)] : index + %c0_offset = index.constant 0 : offset + scf.if %publish_word { + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf16> + %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store %truncated, %output_view[%token, %channel] : vector<4xf16>, view<[%bounded_token_count]x[%bounded_hidden_size]xf16> + } + template.return +} + +template.def<@ggml.publish_f32.q8_1_x4_vector4> device @ggml_publish_f32_q8_1_x4_vector4(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + template.apply<@ggml.quantize_q8_1_x4.publish_vector4>(%publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} + +func.def inline @ggml_publish_f32_row_bytes(%format: index, %hidden_size: index) -> (offset) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c81 = index.constant 81 : index + %c128 = index.constant 128 : index + %zero_bytes = index.constant 0 : offset + %f16_element_bytes = index.constant 2 : offset + %f32_element_bytes = index.constant 4 : offset + %q8_group_bytes = index.constant 144 : offset + %is_f16 = index.cmp eq, %format, %c16 : index + %is_f32 = index.cmp eq, %format, %c32 : index + %is_q8_1 = index.cmp eq, %format, %c81 : index + %f16_row_bytes = index.scale %hidden_size, %f16_element_bytes : index, offset -> offset + %f32_row_bytes = index.scale %hidden_size, %f32_element_bytes : index, offset -> offset + %q8_group_count = index.div %hidden_size, %c128 : index + %q8_row_bytes = index.scale %q8_group_count, %q8_group_bytes : index, offset -> offset + %selected_f16 = scf.select %is_f16, %f16_row_bytes, %zero_bytes : offset + %selected_f32 = scf.select %is_f32, %f32_row_bytes, %selected_f16 : offset + %selected = scf.select %is_q8_1, %q8_row_bytes, %selected_f32 : offset + func.return %selected : offset +} + +template.def<@ggml.publish_f32.next_vector4> device @ggml_publish_f32_next_vector4(%format: index, %publish_word: i1, %token: index, %channel: index, %token_count: index, %hidden_size: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c81 = index.constant 81 : index + %is_f16 = index.cmp eq, %format, %c16 : index + %is_f32 = index.cmp eq, %format, %c32 : index + %is_q8_1 = index.cmp eq, %format, %c81 : index + %publish_f16 = scalar.andi %publish_word, %is_f16 : i1 + %publish_f32 = scalar.andi %publish_word, %is_f32 : i1 + template.apply<@ggml.publish_f32.f16_vector4>(%publish_f16, %token, %channel, %token_count, %hidden_size, %values, %output) : (i1, index, index, index, index, vector<4xf32>, buffer) + template.apply<@ggml.publish_f32.f32_vector4>(%publish_f32, %token, %channel, %token_count, %hidden_size, %values, %output) : (i1, index, index, index, index, vector<4xf32>, buffer) + scf.if %is_q8_1 { + %row_bytes = func.call @ggml_publish_f32_row_bytes(%format, %hidden_size) : (index, index) -> (offset) + %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset + template.apply<@ggml.publish_f32.q8_1_x4_vector4>(%publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %output) : (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/q4_k_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q4_k_f16.loom new file mode 100644 index 000000000000..73e457665b28 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q4_k_f16.loom @@ -0,0 +1,100 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Decodes four adjacent values from one 32-value Q4_K group for FP16 matrix +// staging. The group and packet coordinates match the contiguous K dimension +// consumed by WMMA tiles. +func.def inline @ggml_q4k_f16_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %block_bytes = index.constant 144 : offset + %scale_offset = index.constant 4 : offset + %code_offset = index.constant 16 : offset + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_offset : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<3xi32> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %half0 = scalar.fptrunc %value0 : f32 to f16 + %half1 = scalar.fptrunc %value1 : f32 to f16 + %half2 = scalar.fptrunc %value2 : f32 to f16 + %half3 = scalar.fptrunc %value3 : f32 to f16 + %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> + func.return %result : vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom new file mode 100644 index 000000000000..b461ffd1c153 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom @@ -0,0 +1,182 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Decodes four adjacent values from one 32-value Q6_K group for FP16 matrix +// staging. The group and packet coordinates match the contiguous K dimension +// consumed by WMMA tiles. +func.def inline @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %ql_half_word_base = index.mul %half, %c16 : index + %ql_side_word_add = index.mul %ql_side, %c8 : index + %ql_word_base = index.add %ql_half_word_base, %ql_side_word_add : index + %ql_word_index = index.add %ql_word_base, %bounded_packet : index + %qh_half_word_base = index.mul %half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_group_base = index.mul %bounded_group, %c2 : index + %scale_index = index.add %scale_group_base, %scale_packet_half : index + %ql_word = vector.load %ql_view[%ql_word_index] : view<32xi32> -> vector<1xi32> + %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 + %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} + +// Decode adjacent K32 groups using their shared high bits and scales. +func.def public inline @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) { + %padding = vector.constant 0 : vector<2xi32> + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %block_bytes = index.constant 210 : offset + %qh_byte_add = index.constant 128 : offset + %scale_byte_add = index.constant 192 : offset + %d_byte_add = index.constant 208 : offset + %c4_i32v = vector.constant 4 : vector<4xi32> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %c32_f32v = vector.constant 32.0 : vector<16xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 6), mul(%q6_group, 2)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 6), mul(%packet, 2)] : index + %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset + %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset + %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset + %d_byte_base = index.add %block_byte_base, %d_byte_add : offset + %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> + %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<4xi32> + %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_half_word_base = index.mul %half, %c16 : index + %ql_word_index0 = index.add %ql_half_word_base, %bounded_packet : index + %ql_word_index1 = index.add %ql_word_index0, %c8 : index + %qh_half_word_base = index.mul %half, %c8 : index + %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> + %qh_shift_index0 = index.mul %group_in_half, %c2 : index + %qh_shift_index1 = index.add %qh_shift_index0, %c2 : index + %qh_shift_i32_0 = index.cast %qh_shift_index0 : index to i32 + %qh_shift_i32_1 = index.cast %qh_shift_index1 : index to i32 + %qh_shift0 = vector.splat %qh_shift_i32_0 : vector<4xi32> + %qh_shift1 = vector.splat %qh_shift_i32_1 : vector<4xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %scale_word_index = index.div %bounded_group, %c2 : index + %scale_shift_index0 = index.mul %scale_packet_half, %c8 : index + %scale_shift_index1 = index.add %scale_shift_index0, %c16 : index + %scale_shift0 = index.cast %scale_shift_index0 : index to i32 + %scale_shift1 = index.cast %scale_shift_index1 : index to i32 + %ql_word0 = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %ql_view[%ql_word_index0] : view<32xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %ql_view[%ql_word_index0] : view<32xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %ql_word1 = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %ql_view[%ql_word_index1] : view<32xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %ql_view[%ql_word_index1] : view<32xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %qh_word = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %scale_word = view.load %scale_view[%scale_word_index] : view<4xi32> -> i32 + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %ql_shifted0 = vector.shrui %ql_word0, %nibble_shift : vector<4xi32> + %ql0 = vector.andi %ql_shifted0, %nibble_mask : vector<4xi32> + %qh_shifted0 = vector.shrui %qh_word, %qh_shift0 : vector<4xi32> + %qh_low0 = vector.andi %qh_shifted0, %high_mask : vector<4xi32> + %qh0 = vector.shli %qh_low0, %c4_i32v : vector<4xi32> + %code0 = vector.ori %ql0, %qh0 : vector<4xi32> + %code_i8_0 = vector.bitcast %code0 : vector<4xi32> to vector<16xi8> + %code_f32_0 = vector.uitofp %code_i8_0 : vector<16xi8> to vector<16xf32> + %centered0 = vector.subf %code_f32_0, %c32_f32v : vector<16xf32> + %scale_shifted0 = scalar.shrui %scale_word, %scale_shift0 : i32 + %scale_i8_0 = scalar.trunci %scale_shifted0 : i32 to i8 + %scale0 = scalar.sitofp %scale_i8_0 : i8 to f32 + %combined_scale0 = scalar.mulf %scale0, %d : f32 + %combined_scale_vector0 = vector.splat %combined_scale0 : vector<16xf32> + %values_f32_0 = vector.mulf %centered0, %combined_scale_vector0 : vector<16xf32> + %values0 = vector.fptrunc %values_f32_0 : vector<16xf32> to vector<16xf16> + %ql_shifted1 = vector.shrui %ql_word1, %nibble_shift : vector<4xi32> + %ql1 = vector.andi %ql_shifted1, %nibble_mask : vector<4xi32> + %qh_shifted1 = vector.shrui %qh_word, %qh_shift1 : vector<4xi32> + %qh_low1 = vector.andi %qh_shifted1, %high_mask : vector<4xi32> + %qh1 = vector.shli %qh_low1, %c4_i32v : vector<4xi32> + %code1 = vector.ori %ql1, %qh1 : vector<4xi32> + %code_i8_1 = vector.bitcast %code1 : vector<4xi32> to vector<16xi8> + %code_f32_1 = vector.uitofp %code_i8_1 : vector<16xi8> to vector<16xf32> + %centered1 = vector.subf %code_f32_1, %c32_f32v : vector<16xf32> + %scale_shifted1 = scalar.shrui %scale_word, %scale_shift1 : i32 + %scale_i8_1 = scalar.trunci %scale_shifted1 : i32 to i8 + %scale1 = scalar.sitofp %scale_i8_1 : i8 to f32 + %combined_scale1 = scalar.mulf %scale1, %d : f32 + %combined_scale_vector1 = vector.splat %combined_scale1 : vector<16xf32> + %values_f32_1 = vector.mulf %centered1, %combined_scale_vector1 : vector<16xf32> + %values1 = vector.fptrunc %values_f32_1 : vector<16xf32> to vector<16xf16> + func.return %values0, %values1 : vector<16xf16>, vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q8_0_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q8_0_f16.loom new file mode 100644 index 000000000000..26f7760ac4b9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q8_0_f16.loom @@ -0,0 +1,26 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Decodes four adjacent values from one 32-value Q8_0 block for FP16 matrix +// staging. +func.def inline @ggml_q8_0_f16_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf16>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 34 : offset + %code_offset = index.constant 2 : offset + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %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<32xi8> + %packet_base = index.mul %bounded_packet, %c4 : index + %codes = vector.load %code_view[%packet_base] : view<32xi8> -> vector<4xi8> + %codes_f32 = vector.sitofp %codes : 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> + %values_f32 = vector.mulf %codes_f32, %d_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q8_1_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q8_1_f16.loom new file mode 100644 index 000000000000..1ccd1fdf3b36 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q8_1_f16.loom @@ -0,0 +1,26 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Decodes four adjacent values from one 32-value Q8_1 block for FP16 matrix +// staging. The block sum field is not part of the row values. +func.def inline @ggml_q8_1_f16_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf16>) { + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 36 : offset + %code_offset = index.constant 4 : offset + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %block_byte_add = index.scale %q8_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %code_byte_base = index.add %block_byte_base, %code_offset : offset + %ds_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi8> + %packet_base = index.mul %bounded_packet, %c4 : index + %codes = vector.load %code_view[%packet_base] : view<32xi8> -> vector<4xi8> + %codes_f32 = vector.sitofp %codes : vector<4xi8> to vector<4xf32> + %d_f16 = view.load %ds_view[%c0] : view<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %d_vector = vector.splat %d : vector<4xf32> + %values_f32 = vector.mulf %codes_f32, %d_vector : vector<4xf32> + %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> + func.return %values : vector<4xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4.loom new file mode 100644 index 000000000000..2866f5671e86 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4.loom @@ -0,0 +1,120 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Publishes four F32 values into GGML's block_q8_1_x4 layout. The caller must +// execute complete eight-lane cohorts with four consecutive channels per lane. +template.decl @ggml.quantize_q8_1_x4.publish_vector4.body(%strict_scaling: i1, %publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.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) + +template.def<@ggml.quantize_q8_1_x4.publish_vector4> device @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) { + %strict_scaling = scalar.constant false : i1 + template.apply<@ggml.quantize_q8_1_x4.publish_vector4.body>(%strict_scaling, %publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %output) : (i1, i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} + +template.decl @ggml.quantize_q8_1_x4.publish_vector4_strict(%publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.def<@ggml.quantize_q8_1_x4.publish_vector4_strict> device @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) { + %strict_scaling = scalar.constant true : i1 + template.apply<@ggml.quantize_q8_1_x4.publish_vector4.body>(%strict_scaling, %publish_word, %token_output_byte_base, %channel, %values, %scratch_values, %scratch_d, %output) : (i1, i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} + +template.def<@ggml.quantize_q8_1_x4.publish_vector4.body> device @ggml_quantize_q8_1_x4_publish_vector4_body(%strict_scaling: i1, %publish_word: i1, %token_output_byte_base: offset, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) { + %workitem = kernel.workitem.id : 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 + %c32 = index.constant 32 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %word_in_block = index.rem %workitem, %c8 : index + %absolute_values = vector.absf %values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %shuffle_width = scalar.constant 8 : i32 + %xor1 = scalar.constant 1 : i32 + %max_peer1, %max_valid1 = kernel.subgroup.shuffle %thread_max, %xor1, %shuffle_width : f32, i32, i32 + %max1 = scalar.maxnumf %thread_max, %max_peer1 : f32 + %xor2 = scalar.constant 2 : i32 + %max_peer2, %max_valid2 = kernel.subgroup.shuffle %max1, %xor2, %shuffle_width : f32, i32, i32 + %max2 = scalar.maxnumf %max1, %max_peer2 : f32 + %xor4 = scalar.constant 4 : i32 + %max_peer4, %max_valid4 = kernel.subgroup.shuffle %max2, %xor4, %shuffle_width : f32, i32, i32 + %max4 = scalar.maxnumf %max2, %max_peer4 : f32 + %d0 = scf.if %strict_scaling -> (f32) { + %precise = scalar.divf %max4, %c127 : f32 + scf.yield %precise : f32 + } else { + %fast = scalar.divf %max4, %c127 : f32 + scf.yield %fast : f32 + } + %d = scf.select %publish_word, %d0, %c0_f32 : f32 + %is_block_leader = index.cmp eq, %word_in_block, %c0 : index + %writes_block_d = scalar.andi %publish_word, %is_block_leader : i1 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scf.if %strict_scaling -> (f32) { + %precise = scalar.divf %c1_f32, %d : f32 + scf.yield %precise : f32 + } else { + %fast = scalar.divf %c1_f32, %d : f32 + scf.yield %fast : f32 + } + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + scf.if %publish_word { + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_qs = buffer.view %output[%payload_byte_offset] : buffer -> view<32xi32> + %block_word_base = index.mul %block_in_group, %c8 : index + %packed_word_index0 = index.add %block_word_base, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + %sum_peer1, %sum_valid1 = kernel.subgroup.shuffle %thread_quantized_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum1 = scalar.addf %thread_quantized_sum, %sum_peer1 : f32 + %sum_peer2, %sum_valid2 = kernel.subgroup.shuffle %sum1, %xor2, %shuffle_width : f32, i32, i32 + %sum2 = scalar.addf %sum1, %sum_peer2 : f32 + %sum_peer4, %sum_valid4 = kernel.subgroup.shuffle %sum2, %xor4, %shuffle_width : f32, i32, i32 + %quantized_sum = scalar.addf %sum2, %sum_peer4 : f32 + scf.if %writes_block_d { + %s = scf.if %strict_scaling -> (f32) { + %precise = scalar.mulf %quantized_sum, %d : f32 + scf.yield %precise : f32 + } else { + %fast = scalar.mulf %quantized_sum, %d : f32 + scf.yield %fast : f32 + } + %q8_block = index.div %channel, %c32 : index + %physical_group = index.div %q8_block, %c4 : index + %block_in_group = index.rem %q8_block, %c4 : index + %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset + %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset + %group_ds = buffer.view %output[%group_byte_offset] : buffer -> view<8xf16> + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + %s_index = index.add %ds_index, %c1 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4_group.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4_group.loom new file mode 100644 index 000000000000..02f6eff75c9e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/quantize_q8_1_x4_group.loom @@ -0,0 +1,157 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.quantize_q8_1_x4.group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) + +config.decl @ggml.quantize_q8_1_x4.group_capacity : %value: index where [range(%value, 1, 524288)] + +// Loads one logical 32-element Q8_1 block from the four-way physical packing. +// `row_byte_base` addresses the first x4 group for one activation row. +func.def inline @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %s_index = index.add %d_index, %c1 : index + %word_index = index.mul %inner, %c8 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %s_f16 = view.load %ds_view[%s_index] : view<8xf16> -> f16 + %words = vector.load %payload_view[%word_index] : view<32xi32> -> vector<8xi32> + %values = vector.bitcast %words : vector<8xi32> to vector<32xi8> + %d = scalar.extf %d_f16 : f16 to f32 + %s = scalar.extf %s_f16 : f16 to f32 + func.return %values, %d, %s : vector<32xi8>, f32, f32 +} + +// Loads one packed four-value word and its block scale. Dot-product kernels +// use this narrower form so a lane does not load the other seven words owned +// by its subgroup peers. +func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %group = index.div %q8_block, %c4 : index + %inner0 = index.rem %q8_block, %c4 : index + %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index + %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index + %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset + %group_byte_base = index.add %row_byte_base, %group_byte_add : offset + %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset + %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> + %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> + %d_index = index.mul %inner, %c2 : index + %inner_word_base = index.mul %inner, %c8 : index + %word_index = index.add %inner_word_base, %word0 : index + %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 + %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> + %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> + %d = scalar.extf %d_f16 : f16 to f32 + func.return %values, %d : vector<4xi8>, f32 +} + +// Packs one explicit physical group. Callers provide the complete physical +// group domain and an in-range ordinal, and guarantee that one complete +// 32-lane wave executes after all 128 source values are visible. The uniform +// publication predicate suppresses only destination writes so every lane still +// reaches the workgroup barriers. +template.def<@ggml.quantize_q8_1_x4.group_body> device @ggml_quantize_q8_1_x4_group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane = kernel.workitem.id : 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 + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 128 : offset + %scratch_bytes = index.constant 144 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %block_in_group0 = index.div %lane, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %group_element_base = index.mul %group, %c128 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %block_word_add = index.mul %block_in_group, %c8 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_block_base = index.add %group_element_base, %block_element_add : index + %input_index = index.add %input_block_base, %word_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> + view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index + scf.if %is_cohort_leader { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %packed_word_index0 = index.add %block_word_add, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + scf.if %publish_output { + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 + scf.if %publishes_metadata { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + %s_index = index.add %ds_index, %c1 : index + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + template.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 new file mode 100644 index 000000000000..8f12fb9d16e1 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rope/rope_f32.loom @@ -0,0 +1,347 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +func.def inline @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %position_vector = vector.splat %position : vector<2xf32> + %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 + %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> + %scaled_theta = vector.divf %theta, %freq_factors : vector<2xf32> + %angles = vector.mulf %position_vector, %scaled_theta : vector<2xf32> + %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> + %cosines = vector.costurnsf %turns : vector<2xf32> + %sines = vector.sinturnsf %turns : vector<2xf32> + %x_cosines = vector.mulf %x_values, %cosines : vector<2xf32> + %y_sines = vector.mulf %y_values, %sines : vector<2xf32> + %x_sines = vector.mulf %x_values, %sines : vector<2xf32> + %y_cosines = vector.mulf %y_values, %cosines : vector<2xf32> + %rotated_x = vector.subf %x_cosines, %y_sines : vector<2xf32> + %rotated_y = vector.addf %x_sines, %y_cosines : vector<2xf32> + func.return %rotated_x, %rotated_y : vector<2xf32>, vector<2xf32> +} + +template.decl @ggml.rope_f32.rotate_scalar_pair(%x: f32, %y: f32, %cosine: f32, %sine: f32) -> (f32, f32) + +template.def<@ggml.rope_f32.rotate_scalar_pair> device @ggml_rope_f32_rotate_scalar_pair(%x: f32, %y: f32, %cosine: f32, %sine: f32) -> (f32, f32) { + %x_cosine = scalar.mulf %x, %cosine : f32 + %y_sine = scalar.mulf %y, %sine : f32 + %rotated_x = scalar.subf %x_cosine, %y_sine : f32 + %x_sine = scalar.mulf %x, %sine : f32 + %y_cosine = scalar.mulf %y, %cosine : f32 + %rotated_y = scalar.addf %x_sine, %y_cosine : f32 + template.return %rotated_x, %rotated_y : f32, f32 +} + +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, %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) + +template.decl @ggml.rope_f32.rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +template.decl @ggml.rope_f32.rotate_normal_packet(%position: f32, %input_base: index, %normal_channel: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<4xf32>) + +template.decl @ggml.rope_f32.copy_packet(%input_base: index, %normal_channel: index, %bounded_input_span: index, %input: buffer) -> (vector<4xf32>) + +template.decl @ggml.rows.cache_publish.safe_i64_row(%arg0: index, %arg1: index, %arg2: index, %arg3: buffer) -> (i1, index) + +template.decl @ggml.rows.cache_publish.f32_vector2(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<2xf32>, %arg7: buffer) + +template.decl @ggml.rows.cache_publish.f32_vector4(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) + +config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.rope_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_f32.input_span : %value: index where [range(%value, 1, 1073741824)] + +config.decl @ggml.rope_f32.mscale : f32 + +config.decl @ggml.rope_f32.mode : %value: index where [range(%value, 0, 2)] + +config.decl @ggml.rope_scale_f32.scale : f32 + +config.decl @ggml.rope_set_rows_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_set_rows_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @ggml.rope_set_rows_f32.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.rope_set_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.rope_set_rows_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_set_rows_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] + +config.decl @ggml.rope_set_rows_f32.input_span : %value: index where [range(%value, 1, 1073741824)] + +config.decl @ggml.rope_set_rows_f32.mscale : f32 + +config.decl @ggml.rope_set_rows_f32.output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.rope_set_rows_f32.mode : %value: index where [range(%value, 0, 2)] + +template.def<@ggml.rope_f32.rotate_neox_packet> device @ggml_rope_f32_rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { + %c2 = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %mscale_vector = vector.splat %mscale : vector<2xf32> + %paired_channel = index.add %pair_channel, %half_n_dims : index + %low_index0 = index.add %input_base, %pair_channel : index + %high_index0 = index.add %input_base, %paired_channel : index + %low_end0 = index.add %low_index0, %c2 : index + %high_end0 = index.add %high_index0, %c2 : index + %low_index, %low_end = index.assume %low_index0, %low_end0 [le(%low_end0, %bounded_input_span)] : index, index + %high_index, %high_end = index.assume %high_index0, %high_end0 [le(%high_end0, %bounded_input_span)] : index, index + %low_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%low_end]xf32> + %high_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%high_end]xf32> + %low_values = vector.load %low_input_view[%low_index] : view<[%low_end]xf32> -> vector<2xf32> + %high_values = vector.load %high_input_view[%high_index] : view<[%high_end]xf32> -> vector<2xf32> + %rotated_low, %rotated_high = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %low_values, %high_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %scaled_low = vector.mulf %rotated_low, %mscale_vector : vector<2xf32> + %scaled_high = vector.mulf %rotated_high, %mscale_vector : vector<2xf32> + template.return %scaled_low, %scaled_high : vector<2xf32>, vector<2xf32> +} + +template.def<@ggml.rope_f32.rotate_normal_packet> device @ggml_rope_f32_rotate_normal_packet(%position: f32, %input_base: index, %normal_channel: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<4xf32>) { + %c4 = index.constant 4 : index + %c0_offset = index.constant 0 : offset + %mscale_vector = vector.splat %mscale : vector<2xf32> + %normal_index0 = index.add %input_base, %normal_channel : index + %normal_end0 = index.add %normal_index0, %c4 : index + %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index + %normal_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%normal_end]xf32> + %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> + %x0 = vector.extract %values[0] : vector<4xf32> -> f32 + %y0 = vector.extract %values[1] : vector<4xf32> -> f32 + %x1 = vector.extract %values[2] : vector<4xf32> -> f32 + %y1 = vector.extract %values[3] : vector<4xf32> -> f32 + %x_values = vector.from_elements %x0, %x1 : vector<2xf32> + %y_values = vector.from_elements %y0, %y1 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %scaled_x = vector.mulf %rotated_x, %mscale_vector : vector<2xf32> + %scaled_y = vector.mulf %rotated_y, %mscale_vector : vector<2xf32> + %scaled_x0 = vector.extract %scaled_x[0] : vector<2xf32> -> f32 + %scaled_x1 = vector.extract %scaled_x[1] : vector<2xf32> -> f32 + %scaled_y0 = vector.extract %scaled_y[0] : vector<2xf32> -> f32 + %scaled_y1 = vector.extract %scaled_y[1] : vector<2xf32> -> f32 + %rotated = vector.from_elements %scaled_x0, %scaled_y0, %scaled_x1, %scaled_y1 : vector<4xf32> + template.return %rotated : vector<4xf32> +} + +template.def<@ggml.rope_f32.copy_packet> device @ggml_rope_f32_copy_packet(%input_base: index, %normal_channel: index, %bounded_input_span: index, %input: buffer) -> (vector<4xf32>) { + %c4 = index.constant 4 : index + %c0_offset = index.constant 0 : offset + %normal_index0 = index.add %input_base, %normal_channel : index + %normal_end0 = index.add %normal_index0, %c4 : index + %normal_index, %normal_end = index.assume %normal_index0, %normal_end0 [le(%normal_end0, %bounded_input_span)] : index, index + %normal_input_view = buffer.view %input[%c0_offset] : buffer -> view<[%normal_end]xf32> + %values = vector.load %normal_input_view[%normal_index] : view<[%normal_end]xf32> -> vector<4xf32> + template.return %values : vector<4xf32> +} + +template.def<@ggml.rope_f32.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 + %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_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 + %token0 = kernel.workgroup.id : index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %mode_neox = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %output_scale2 = vector.splat %output_scale : vector<2xf32> + %output_scale4 = vector.splat %output_scale : vector<4xf32> + %half_n_dims = index.div %n_dims, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %rotary_packet_count = index.div %n_dims, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %rotary_channel = index.cmp ult, %channel0, %rotary_packet_count : index + %tail_channel = index.cmp uge, %channel0, %rotary_packet_count : index + %copy_channel = scalar.andi %active_channel, %tail_channel : i1 + %publish = scalar.andi %valid_token, %active_channel : i1 + %publish_rotary0 = scalar.andi %valid_token, %rotary_channel : i1 + %publish_rotary = scalar.andi %publish_rotary0, %active_channel : i1 + %publish_copy = scalar.andi %valid_token, %copy_channel : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %is_neox = index.cmp eq, %mode, %mode_neox : index + %is_normal = index.cmp eq, %mode, %c0 : index + %positions_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %input, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer, buffer + scf.if %publish_rotary { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %normal_channel = index.mul %channel, %c4 : index + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + scf.if %is_neox { + %paired_channel = index.add %pair_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %output_low = vector.mulf %scaled_low, %output_scale2 : vector<2xf32> + %output_high = vector.mulf %scaled_high, %output_scale2 : vector<2xf32> + vector.store %output_low, %output_view[%token, %head, %pair_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + vector.store %output_high, %output_view[%token, %head, %paired_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + } + scf.if %is_normal { + %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %input_base, %normal_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) + %output_rotated = vector.mulf %rotated, %output_scale4 : vector<4xf32> + vector.store %output_rotated, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + } + } + scf.if %publish_copy { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %normal_channel = index.mul %channel, %c4 : index + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %normal_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + %output_values = vector.mulf %values, %output_scale4 : vector<4xf32> + vector.store %output_values, %output_view[%token, %head, %normal_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%head_size]xf32> + } + template.return +} + +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, %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 + %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 +} + +template.def<@ggml.rope_token_scale_f32.body> device @ggml_rope_token_scale_f32_body(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %token0 = kernel.workgroup.id : index + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %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 + %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 +} + +template.def<@ggml.rope_set_rows_f32.body> device @ggml_rope_set_rows_f32_body(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index + %head_count0 = config.get @ggml.rope_set_rows_f32.head_count : index + %head_size0 = config.get @ggml.rope_set_rows_f32.head_size : index + %n_dims0 = config.get @ggml.rope_set_rows_f32.n_dims : index + %input_stride1_0 = config.get @ggml.rope_set_rows_f32.input_stride1 : index + %input_stride2_0 = config.get @ggml.rope_set_rows_f32.input_stride2 : index + %input_span0 = config.get @ggml.rope_set_rows_f32.input_span : index + %mscale = config.get @ggml.rope_set_rows_f32.mscale : f32 + %output_format = config.get @ggml.rope_set_rows_f32.output_format : index + %mode = config.get @ggml.rope_set_rows_f32.mode : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index + %head_count, %head_size, %n_dims = index.assume %head_count0, %head_size0, %n_dims0 [range(%head_count0, 1, 64), range(%head_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%head_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %head_size0)] : index, index, index + %input_stride1, %input_stride2 = index.assume %input_stride1_0, %input_stride2_0 [range(%input_stride1_0, 4, 1073741824), range(%input_stride2_0, 4, 1073741824)] : index, index + %head = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %channel0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %mode_neox = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %half_n_dims = index.div %n_dims, %c2 : index + %pair_packet_count = index.div %head_size, %c4 : index + %rotary_packet_count = index.div %n_dims, %c4 : index + %cache_hidden_size = index.mul %head_count, %head_size : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %active_channel = index.cmp ult, %channel0, %pair_packet_count : index + %rotary_channel = index.cmp ult, %channel0, %rotary_packet_count : index + %tail_channel = index.cmp uge, %channel0, %rotary_packet_count : index + %copy_channel = scalar.andi %active_channel, %tail_channel : i1 + %publish0 = scalar.andi %valid_token, %active_channel : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %is_neox = index.cmp eq, %mode, %mode_neox : index + %is_normal = index.cmp eq, %mode, %c0 : index + %positions_noalias, %indices_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %cache_noalias = buffer.assume.noalias %positions, %indices, %input, %theta, %freq_factors, %cache : buffer, buffer, buffer, buffer, buffer, buffer + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %valid_row, %cache_row = template.apply<@ggml.rows.cache_publish.safe_i64_row>(%token, %bounded_token_count, %bounded_cache_row_count, %indices_noalias) : (index, index, index, buffer) -> (i1, index) + %publish = scalar.andi %publish0, %valid_row : i1 + %publish_rotary = scalar.andi %publish, %rotary_channel : i1 + %publish_copy = scalar.andi %publish, %copy_channel : i1 + scf.if %publish_rotary { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %pair_channel = index.mul %channel, %c2 : index + %normal_channel = index.mul %channel, %c4 : index + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %cache_head_base = index.mul %head, %head_size : index + scf.if %is_neox { + %paired_channel = index.add %pair_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %low_cache_channel = index.add %cache_head_base, %pair_channel : index + %high_cache_channel = index.add %cache_head_base, %paired_channel : index + template.apply<@ggml.rows.cache_publish.f32_vector2>(%output_format, %publish, %cache_row, %low_cache_channel, %bounded_cache_row_count, %cache_hidden_size, %scaled_low, %cache_noalias) : (index, i1, index, index, index, index, vector<2xf32>, buffer) + template.apply<@ggml.rows.cache_publish.f32_vector2>(%output_format, %publish, %cache_row, %high_cache_channel, %bounded_cache_row_count, %cache_hidden_size, %scaled_high, %cache_noalias) : (index, i1, index, index, index, index, vector<2xf32>, buffer) + } + scf.if %is_normal { + %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %input_base, %normal_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) + %cache_channel = index.add %cache_head_base, %normal_channel : index + template.apply<@ggml.rows.cache_publish.f32_vector4>(%output_format, %publish, %cache_row, %cache_channel, %bounded_cache_row_count, %cache_hidden_size, %rotated, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) + } + } + scf.if %publish_copy { + %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index + %normal_channel = index.mul %channel, %c4 : index + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %normal_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + %cache_head_base = index.mul %head, %head_size : index + %cache_channel = index.add %cache_head_base, %normal_channel : index + template.apply<@ggml.rows.cache_publish.f32_vector4>(%output_format, %publish, %cache_row, %cache_channel, %bounded_cache_row_count, %cache_hidden_size, %values, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_fused.loom new file mode 100644 index 000000000000..1d95f9fae09b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_fused.loom @@ -0,0 +1,91 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.router.fused.experts_per_wave() -> (index) + +template.decl @qwen3_moe.router.fused.projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) + +amdgpu.target @qwen3_moe_router_fused_gfx1151 + +template.def<@qwen3_moe.router.fused.experts_per_wave> target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_experts_per_wave() -> (index) { + %c2 = index.constant 2 : index + template.return %c2 : index +} + +template.def<@qwen3_moe.router.fused.experts_per_wave> requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_experts_per_wave() -> (index) { + %c4 = index.constant 4 : index + template.return %c4 : index +} + +template.def<@qwen3_moe.router.fused.projection> device target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %expert_base = index.mul %expert_tile, %c2 : index + %expert1 = index.add %expert_base, %c1 : index + %lane_channel = index.mul %lane, %c4 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> + %c0_f32x2 = vector.splat %c0_f32 : vector<2xf32> + %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<2xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<2xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1 : vector<2xf32> + scf.yield %next_accumulators : vector<2xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> + scf.if %writes_projection { + vector.store %sums, %logits_view[%token, %expert_base] : vector<2xf32>, view<[%token_count]x[%expert_count]xf32> + } + template.return +} + +template.def<@qwen3_moe.router.fused.projection> device requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { + %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 + %c256 = index.constant 256 : index + %c0_f32 = scalar.constant 0.0 : f32 + %expert_base = index.mul %expert_tile, %c4 : index + %expert1 = index.add %expert_base, %c1 : index + %expert2 = index.add %expert_base, %c2 : index + %expert3 = index.add %expert_base, %c3 : index + %lane_channel = index.mul %lane, %c4 : index + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> + %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> + %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> + %c0_f32x4 = vector.splat %c0_f32 : vector<4xf32> + %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%c4) schedule(interleaved) { + %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values2 = vector.load %weight_view[%expert2, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %weight_values3 = vector.load %weight_view[%expert3, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> + %accumulator0 = vector.extract %accumulators[0] : vector<4xf32> -> f32 + %accumulator1 = vector.extract %accumulators[1] : vector<4xf32> -> f32 + %accumulator2 = vector.extract %accumulators[2] : vector<4xf32> -> f32 + %accumulator3 = vector.extract %accumulators[3] : vector<4xf32> -> f32 + %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 + %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 + %next2 = vector.dotf %input_values, %weight_values2, %accumulator2 : vector<4xf32>, vector<4xf32>, f32 + %next3 = vector.dotf %input_values, %weight_values3, %accumulator3 : vector<4xf32>, vector<4xf32>, f32 + %next_accumulators = vector.from_elements %next0, %next1, %next2, %next3 : vector<4xf32> + scf.yield %next_accumulators : vector<4xf32> + } + %sums = kernel.subgroup.reduce %lane_sums : vector<4xf32> + scf.if %writes_projection { + vector.store %sums, %logits_view[%token, %expert_base] : vector<4xf32>, view<[%token_count]x[%expert_count]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_storage.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_storage.loom new file mode 100644 index 000000000000..39cd1a9bea21 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_projection_storage.loom @@ -0,0 +1,19 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.router_projection.storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) + +amdgpu.target @qwen3_moe_router_projection_gfx1151 + +// Selects the target's preferred storage contract without exposing addressing +// policy at the projection call site. +template.def<@qwen3_moe.router_projection.storage> target(@qwen3_moe_router_projection_gfx1151) priority(20) @qwen3_moe_router_projection_descriptor_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { + %input_descriptor = buffer.assume.memory_space %input : buffer + %weight_descriptor = buffer.assume.memory_space %weight : buffer + %output_descriptor = buffer.assume.memory_space %output : buffer + template.return %input_descriptor, %weight_descriptor, %output_descriptor : buffer, buffer, buffer +} + +template.def<@qwen3_moe.router_projection.storage> priority(1) @qwen3_moe_router_projection_global_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { + template.return %input, %weight, %output : buffer, buffer, buffer +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom new file mode 100644 index 000000000000..45aadcd71549 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/router_topk_row.loom @@ -0,0 +1,100 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.router.top8.row(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) + +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] + +config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] + +// Selects and normalizes one logical router row. The caller owns mapping a +// subgroup to a safe physical row and tells the helper whether that row should +// publish. Keeping row semantics independent of launch geometry lets fused +// producers consume this exact tie-breaking and normalization contract. +template.def<@qwen3_moe.router.top8.row> device requires [#target.subgroup.size<64>] @qwen3_moe_router_top8_row_f32(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %expert_count0 = config.get @qwen3_moe.router.expert_count : index + %route_count0 = config.get @qwen3_moe.router.route_count : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_route_id_stride0 = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %expert_count, %route_count, %bounded_route_id_stride = index.assume %expert_count0, %route_count0, %bounded_route_id_stride0 [range(%expert_count0, 32, 512), mul(%expert_count0, 32), range(%route_count0, 1, 32), le(%route_count0, %expert_count0), le(%route_count0, %bounded_route_id_stride0)] : index, index, index + %safe_token, %launch_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %wave_size = target.subgroup.size : index + %negative_large = scalar.constant -3.4028234663852886e+38 : f32 + %largest_i32 = scalar.constant 2147483647 : i32 + %c0_offset = index.constant 0 : offset + %experts_per_lane0 = index.div %expert_count, %wave_size : index + %experts_per_lane = index.assume %experts_per_lane0 [range(%experts_per_lane0, 1, 16)] : index + %lane_expert_base = index.mul %lane, %experts_per_lane : index + %route_id_storage_count = index.mul %launch_token_count, %bounded_route_id_stride : index + %route_weight_storage_count = index.mul %launch_token_count, %route_count : index + %logits_noalias, %route_ids_noalias, %route_weights_noalias = buffer.assume.noalias %logits, %route_ids, %route_weights : buffer, buffer, buffer + %logits_view = buffer.view %logits_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> + %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%route_id_storage_count]xi32> + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%route_weight_storage_count]xf32> + %initial_logits = vector.load %logits_view[%safe_token, %lane_expert_base] : view<[%launch_token_count]x[%expert_count]xf32> -> vector<[%experts_per_lane]xf32> + %remaining_final, %selected_logit = scf.for %route = [%c0 to %route_count step %c1](%remaining_logits = %initial_logits : vector<[%experts_per_lane]xf32>, %lane_selected_logit = %negative_large : f32) -> (vector<[%experts_per_lane]xf32>, f32) { + %local_value, %local_id = scf.for %slot = [%c0 to %experts_per_lane step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { + %candidate_value = vector.extract %remaining_logits[%slot] : vector<[%experts_per_lane]xf32> -> f32 + %candidate_expert = index.add %lane_expert_base, %slot : index + %candidate_id = index.cast %candidate_expert : index to i32 + %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 + %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 + %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 + %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 + %is_better = scalar.ori %is_greater, %is_lower_tie : i1 + %next_value, %next_id = scf.if %is_better -> (f32, i32) { + scf.yield %candidate_value, %candidate_id : f32, i32 + } else { + scf.yield %best_value, %best_id : f32, i32 + } + scf.yield %next_value, %next_id : f32, i32 + } + %winner_value = kernel.subgroup.reduce %local_value : f32 + %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 + %winner_lane_mask = kernel.subgroup.vote.ballot %matches_winner : i1 -> i64 + %nonzero_winner_lane_mask = scalar.assume %winner_lane_mask [ne(%winner_lane_mask, 0)] : i64 + %winner_lane_i64 = scalar.cttzi %nonzero_winner_lane_mask : i64 + %winner_lane0 = index.cast %winner_lane_i64 : i64 to index + %winner_lane = index.assume %winner_lane0 [range(%winner_lane0, 0, 63)] : index + %local_expert0 = index.cast %local_id : i32 to index + %local_expert = index.assume %local_expert0 [range(%local_expert0, 0, 511)] : index + %winner_slot = index.rem %local_expert, %experts_per_lane : index + %owns_winner = index.cmp eq, %lane, %winner_lane : index + %next_remaining_logits = scf.if %owns_winner -> (vector<[%experts_per_lane]xf32>) { + %removed = vector.insert %negative_large into %remaining_logits[%winner_slot] : f32, vector<[%experts_per_lane]xf32> + scf.yield %removed : vector<[%experts_per_lane]xf32> + } else { + scf.yield %remaining_logits : vector<[%experts_per_lane]xf32> + } + %publishes_route_id = scalar.andi %valid_token, %owns_winner : i1 + scf.if %publishes_route_id { + %route_id_token_base = index.mul %safe_token, %bounded_route_id_stride : index + %route_id_index = index.add %route_id_token_base, %route : index + view.store %local_id, %route_ids_view[%route_id_index] : i32, view<[%route_id_storage_count]xi32> + } + %lane_publishes_route = index.cmp eq, %lane, %route : index + %publishes_route = scalar.andi %valid_token, %lane_publishes_route : i1 + %next_lane_selected_logit = scf.if %publishes_route -> (f32) { + scf.yield %winner_value : f32 + } else { + scf.yield %lane_selected_logit : f32 + } + scf.yield %next_remaining_logits, %next_lane_selected_logit : vector<[%experts_per_lane]xf32>, f32 + } + %selected_max = kernel.subgroup.reduce %selected_logit : f32 + %selected_delta = scalar.subf %selected_logit, %selected_max : f32 + %unnormalized_weight = scalar.expf %selected_delta : f32 + %selected_sum = kernel.subgroup.reduce %unnormalized_weight : f32 + %route_weight = scalar.divf %unnormalized_weight, %selected_sum : f32 + %lane_publishes_weight = index.cmp ult, %lane, %route_count : index + %publishes_weight = scalar.andi %valid_token, %lane_publishes_weight : i1 + scf.if %publishes_weight { + %route_weight_token_base = index.mul %safe_token, %route_count : index + %route_weight_index = index.add %route_weight_token_base, %lane : index + view.store %route_weight, %route_weights_view[%route_weight_index] : f32, view<[%route_weight_storage_count]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_f16_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_f16_f32.loom new file mode 100644 index 000000000000..f3f92446428d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_f16_f32.loom @@ -0,0 +1,57 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.routing.weighted_reduce_f16_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +// Preserves route-order FP32 accumulation before loading and adding the +// residual. The caller owns launch geometry and the public ABI. +template.def<@ggml.routing.weighted_reduce_f16_f32> device @ggml_routing_weighted_reduce_f16_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) { + %route_count = config.get @qwen3_moe.routed_down.route_count : index + %output_size = config.get @qwen3_moe.routed_down.output_size : index + %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %assignment_count = index.mul %token_count, %bounded_route_count : index + %channel_tile_base = index.mul %channel_tile, %c256 : index + %lane_channel_add = index.mul %lane, %c4 : index + %channel = index.add %channel_tile_base, %lane_channel_add : index + %active_token = index.cmp ult, %token0, %token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %publishes_output = scalar.andi %active_token, %valid_channel : i1 + %route_weights_noalias, %routed_output_noalias, %residual_input_noalias, %output_noalias = buffer.assume.noalias %route_weights, %routed_output, %residual_input, %output : buffer, buffer, buffer, buffer + %route_weight_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> + scf.if %publishes_output { + %token, %output_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<4xi1> + %weighted_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%sum = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll { + %assignment = index.madd %token, %bounded_route_count, %route : index + %routed = vector.load.mask %routed_output_view[%assignment, %channel], %mask, %c0_f16x4 : view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1>, vector<4xf16> + %wide = vector.extf %routed : vector<4xf16> to vector<4xf32> + %route_weight = view.load %route_weight_view[%token, %route] : view<[%token_count]x[%bounded_route_count]xf32> -> f32 + %route_weight_x4 = vector.splat %route_weight : vector<4xf32> + %weighted = vector.mulf %wide, %route_weight_x4 : vector<4xf32> + %next = vector.addf %sum, %weighted : vector<4xf32> + scf.yield %next : vector<4xf32> + } + %residual = vector.load.mask %residual_input_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> + %result = vector.addf %residual, %weighted_sum : vector<4xf32> + vector.store.mask %result, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + 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 new file mode 100644 index 000000000000..df4defa17958 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom @@ -0,0 +1,148 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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 + +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] + +// Preserves the specialized route-order reduction, residual publication, and +// row RMSNorm sequence. The caller owns launch geometry and the public ABI. +template.def<@ggml.routing.weighted_reduce_next_rmsnorm_f32> device requires [#target.subgroup.size<32>] @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) { + %route_count0 = config.get @qwen3_moe.routed_down.route_count : index + %output_size0 = config.get @qwen3_moe.routed_down.output_size : index + %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index + %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 + %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index + %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %channel = index.mul %workitem, %c8 : index + %scratch_bytes = index.constant 32 : offset + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %active_token = index.cmp ult, %token0, %token_count : index + %safe_token0 = scf.select %active_token, %token0, %c0 : index + %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index + %hidden_size_i32 = index.cast %hidden_size : index to i32 + %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 + %assignment_count = index.mul %launch_token_count, %route_count : index + %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer + %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> + %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> + %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> + %next_projection_input_view = buffer.view %next_projection_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<8xf32> + + // The predicate depends only on the workgroup ID, so every workitem either + // executes the complete barrier-bearing row reduction or skips it together. + scf.if %active_token { + %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { + %assignment = index.madd %token, %route_count, %route : index + %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> + %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> + %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 + %route_weight_x8 = vector.splat %route_weight : vector<8xf32> + %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> + %next = vector.addf %sum, %weighted : vector<8xf32> + scf.yield %next : vector<8xf32> + } + %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> + %result = vector.addf %residual, %weighted_sum : vector<8xf32> + %squares = vector.mulf %result, %result : vector<8xf32> + %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<8xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %scale = view.load %scratch_view[%c0] : view<8xf32> -> f32 + %scale_x8 = vector.splat %scale : vector<8xf32> + %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> + %normalized = vector.mulf %result, %scale_x8 : vector<8xf32> + %next_projection = vector.mulf %normalized, %learned_weight : vector<8xf32> + vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + vector.store %next_projection, %next_projection_input_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> + } + 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/rows/cache_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/cache_publish.loom new file mode 100644 index 000000000000..59eaa9d8c006 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/cache_publish.loom @@ -0,0 +1,69 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rows.cache_publish.safe_i64_row(%token: index, %token_count: index, %cache_row_count: index, %indices: buffer) -> (i1, index) + +template.decl @ggml.rows.cache_publish.f32_vector2(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<2xf32>, %cache: buffer) + +template.decl @ggml.rows.cache_publish.f32_vector4(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<4xf32>, %cache: buffer) + +template.def<@ggml.rows.cache_publish.safe_i64_row> device @ggml_rows_cache_publish_safe_i64_row(%token: index, %token_count: index, %cache_row_count: index, %indices: buffer) -> (i1, index) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %c0 = index.constant 0 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %indices_view = buffer.view %indices[%c0_offset] : buffer -> view<[%bounded_token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%bounded_token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %bounded_cache_row_count : index + %publish = scalar.andi %valid_index, %valid_row : i1 + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + template.return %publish, %cache_row : i1, index +} + +template.def<@ggml.rows.cache_publish.f32_vector2> device @ggml_rows_cache_publish_f32_vector2(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<2xf32>, %cache: buffer) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %is_f16 = index.cmp eq, %output_format, %c16 : index + %is_f32 = index.cmp eq, %output_format, %c32 : index + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf16> + %truncated = vector.fptrunc %values : vector<2xf32> to vector<2xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<2xf16>, view<[%cache_row_count]x[%hidden_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<2xf32>, view<[%cache_row_count]x[%hidden_size]xf32> + } + template.return +} + +template.def<@ggml.rows.cache_publish.f32_vector4> device @ggml_rows_cache_publish_f32_vector4(%output_format: index, %publish: i1, %cache_row: index, %channel: index, %cache_row_count: index, %hidden_size: index, %values: vector<4xf32>, %cache: buffer) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %is_f16 = index.cmp eq, %output_format, %c16 : index + %is_f32 = index.cmp eq, %output_format, %c32 : index + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf16> + %truncated = vector.fptrunc %values : vector<4xf32> to vector<4xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<4xf16>, view<[%cache_row_count]x[%hidden_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache[%c0_offset] : buffer -> view<[%cache_row_count]x[%hidden_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<4xf32>, view<[%cache_row_count]x[%hidden_size]xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/get_rows_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/get_rows_decode.loom new file mode 100644 index 000000000000..42c68ec11535 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/rows/get_rows_decode.loom @@ -0,0 +1,169 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Generic GGML GET_ROWS for contiguous rows decoded through F32. +// +// One workitem publishes four adjacent output channels for one requested row. +// Storage-format-specific row decoding is delegated to the common dequant +// motif so Q1_0, Q3_K, Q4_K, Q5_K, Q6_K, IQ4_XS, Q8_0, Q8_1, F16, BF16, and F32 share the same gather shape. +template.decl @ggml.get_rows_f32.body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.get_rows_f32.launch(%hidden_capacity: index, %token_capacity: index) -> (index, index, index, index) + +template.decl @ggml.get_rows_f32.load_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, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) + +template.decl @ggml.get_rows_f32.load_vector4_common(%weight_format: index, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) + +template.decl @ggml.get_rows_f32.next_body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer, %next_output: buffer) + +template.decl @ggml.publish_f32.f32_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +template.decl @ggml.publish_f32.next_vector4(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: view<256xf32>, %arg8: view<32xf32>, %arg9: buffer) + +config.decl @ggml.get_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.get_rows_f32.hidden_capacity : %value: index where [range(%value, 4, 1073741824), mul(%value, 4)] + +config.decl @ggml.get_rows_f32.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.get_rows_f32.next_format : %value: index where [range(%value, 16, 81)] + +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_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>) + +template.def<@ggml.get_rows_f32.load_vector4> device @ggml_get_rows_f32_load_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, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) { + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %token_id = index.assume %token [range(%token, 0, 262207)] : index + %bounded_token_id, %bounded_row_count = index.assume %token_id, %row_count [lt(%token_id, %row_count)] : index, index + %safe_values = scf.if %publish -> (vector<4xf32>) { + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %hidden_size) : (index, index) -> (offset) + %row_byte_base = index.scale %bounded_token_id, %row_bytes : index, offset -> offset + %quant_block = index.div %channel, %c256 : index + %block_channel = index.rem %channel, %c256 : index + %quant_group = index.div %block_channel, %c32 : index + %group_channel = index.rem %block_channel, %c32 : index + %quant_packet = index.div %group_channel, %c4 : index + %values = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight, %row_byte_base, %hidden_size, %quant_block, %quant_group, %quant_packet, %channel) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + %row_values = scf.select %valid_token_id, %values, %c0_f32x4 : vector<4xf32> + scf.yield %row_values : vector<4xf32> + } else { + scf.yield %c0_f32x4 : vector<4xf32> + } + template.return %safe_values : vector<4xf32> +} + +template.def<@ggml.get_rows_f32.load_vector4_common> device @ggml_get_rows_f32_load_vector4_common(%weight_format: index, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %values = template.apply<@ggml.get_rows_f32.load_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, %publish, %valid_token_id, %token, %row_count, %hidden_size, %channel, %weight) : (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, i1, i1, index, index, index, index, buffer) -> (vector<4xf32>) + template.return %values : vector<4xf32> +} + +template.def<@ggml.get_rows_f32.body> device @ggml_get_rows_f32_body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %weight_format = config.get @ggml.get_rows_f32.weight_format : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_row_count = index.assume %row_count [range(%row_count, 1, 262208)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 4, 1073741824), mul(%hidden_size, 4), le(%hidden_size, %hidden_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %packet_tile_base = index.mul %channel_tile, %c256 : index + %packet0 = index.add %packet_tile_base, %workitem : index + %packet_count = index.div %bounded_hidden_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %valid_packet = index.cmp ult, %packet0, %packet_count : index + %publish = scalar.andi %valid_token, %valid_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %packet = scf.select %valid_packet, %packet0, %c0 : index + %token_ids_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %token_ids, %weight, %output : buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %token_id_raw = view.load %token_ids_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %token_id_nonnegative = scalar.cmpi sge, %token_id_raw, %c0_i32 : i32 + %safe_token_id_i32 = scf.select %token_id_nonnegative, %token_id_raw, %c0_i32 : i32 + %token_id = index.cast %safe_token_id_i32 : i32 to index + %channel0 = index.mul %packet, %c4 : index + %channel_end0 = index.add %channel0, %c4 : index + %channel, %channel_end = index.assume %channel0, %channel_end0 [le(%channel_end0, %bounded_hidden_size)] : index, index + %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>) + %values = template.apply<@ggml.get_rows_f32.load_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, %publish, %token_id_nonnegative, %token_id, %bounded_row_count, %bounded_hidden_size, %channel, %weight_noalias) : (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, i1, i1, index, index, index, index, buffer) -> (vector<4xf32>) + scf.if %publish { + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + vector.store %values, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + } + template.return +} + +template.def<@ggml.get_rows_f32.next_body> device @ggml_get_rows_f32_next_body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer, %next_output: buffer) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %weight_format = config.get @ggml.get_rows_f32.weight_format : index + %next_format = config.get @ggml.get_rows_f32.next_format : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_row_count = index.assume %row_count [range(%row_count, 1, 262208)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 256, 1073741824), mul(%hidden_size, 256), le(%hidden_size, %hidden_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %scratch_d_byte_add = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %packet_tile_base = index.mul %channel_tile, %c256 : index + %packet0 = index.add %packet_tile_base, %workitem : index + %packet_count = index.div %bounded_hidden_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %valid_packet = index.cmp ult, %packet0, %packet_count : index + %publish = scalar.andi %valid_token, %valid_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %packet = scf.select %valid_packet, %packet0, %c0 : index + %token_ids_noalias, %weight_noalias, %output_noalias, %next_output_noalias = buffer.assume.noalias %token_ids, %weight, %output, %next_output : buffer, buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %token_id_raw = view.load %token_ids_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %token_id_nonnegative = scalar.cmpi sge, %token_id_raw, %c0_i32 : i32 + %safe_token_id_i32 = scf.select %token_id_nonnegative, %token_id_raw, %c0_i32 : i32 + %token_id = index.cast %safe_token_id_i32 : i32 to index + %channel = index.mul %packet, %c4 : index + %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>) + %values = template.apply<@ggml.get_rows_f32.load_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, %publish, %token_id_nonnegative, %token_id, %bounded_row_count, %bounded_hidden_size, %channel, %weight_noalias) : (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, i1, i1, index, index, index, index, buffer) -> (vector<4xf32>) + template.apply<@ggml.publish_f32.f32_vector4>(%publish, %token, %channel, %bounded_token_count, %bounded_hidden_size, %values, %output_noalias) : (i1, index, index, index, index, vector<4xf32>, buffer) + %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.publish_f32.next_vector4>(%next_format, %publish, %token, %channel, %bounded_token_count, %bounded_hidden_size, %values, %scratch_values, %scratch_d, %next_output_noalias) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.return +} + +template.def<@ggml.get_rows_f32.launch> @ggml_get_rows_f32_launch(%hidden_capacity: index, %token_capacity: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %c256 : index + template.return %packet_tiles, %token_capacity, %c1, %c256 : index, index, index, index +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/dconv4_silu.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/dconv4_silu.loom new file mode 100644 index 000000000000..86b713edf88a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/dconv4_silu.loom @@ -0,0 +1,13 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.dconv4_silu(%x0: f32, %x1: f32, %x2: f32, %x3: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) + +template.def<@llm.ssm_conv.dconv4_silu> device @llm_ssm_conv_dconv4_silu(%x0: f32, %x1: f32, %x2: f32, %x3: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) { + %a0 = scalar.mulf %x0, %w0 : f32 + %a1 = scalar.fmaf %x1, %w1, %a0 : f32 + %a2 = scalar.fmaf %x2, %w2, %a1 : f32 + %a3 = scalar.fmaf %x3, %w3, %a2 : f32 + %result = scalar.siluf %a3 : f32 + template.return %result : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/decode_update.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/decode_update.loom new file mode 100644 index 000000000000..b2f230544136 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/decode_update.loom @@ -0,0 +1,11 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.dconv4_silu(%x0: f32, %x1: f32, %x2: f32, %x3: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) + +template.decl @llm.ssm_conv.decode_update(%prev0: f32, %prev1: f32, %prev2: f32, %new_x: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) + +template.def<@llm.ssm_conv.decode_update> device @llm_ssm_conv_decode_update(%prev0: f32, %prev1: f32, %prev2: f32, %new_x: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) { + %result = template.apply<@llm.ssm_conv.dconv4_silu>(%prev0, %prev1, %prev2, %new_x, %w0, %w1, %w2, %w3) : (f32, f32, f32, f32, f32, f32, f32, f32) -> (f32) + template.return %result : f32 +} 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 new file mode 100644 index 000000000000..6fcf82c6e8ed --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/generic_apply.loom @@ -0,0 +1,39 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.unary_f32.apply(%op: index, %value: f32) -> (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(%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 + %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 + %window_sequence_base = index.mul %sequence, %window_sequence_span : index + %window_channel_base0 = index.mul %channel, %window_rows : index + %window_channel_base = index.assume %window_channel_base0 [range(%window_channel_base0, 0, 34603008)] : index + + %filter_channel_base0 = index.mul %channel, %d_conv : index + %filter_channel_base = index.assume %filter_channel_base0 [range(%filter_channel_base0, 0, 1048560)] : index + + %acc = scf.for %tap = [%c0 to %d_conv step %c1](%running = %c0_f32 : f32) -> (f32) { + %window_row = index.add %token, %tap : index + %window_index0 = index.add %window_sequence_base, %window_channel_base : index + %window_index1 = index.add %window_index0, %window_row : index + %window_index = index.assume %window_index1 [range(%window_index1, 0, 1073741823)] : index + %filter_index0 = index.add %filter_channel_base, %tap : index + %filter_index = index.assume %filter_index0 [range(%filter_index0, 0, 1048575)] : index + %window_value = view.load %window_v[%window_index] : view<1073741824xf32> -> f32 + %filter_value = view.load %filter_v[%filter_index] : view<1048576xf32> -> f32 + %product = scalar.mulf %window_value, %filter_value : f32 + %next = scalar.addf %running, %product : f32 + scf.yield %next : f32 + } + %activated = template.apply<@ggml.unary_f32.apply>(%unary_op, %acc) : (index, f32) -> (f32) + template.return %activated : f32 +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/prefill_finish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/prefill_finish.loom new file mode 100644 index 000000000000..4605180d9448 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/prefill_finish.loom @@ -0,0 +1,55 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.dconv4_silu(%x0: f32, %x1: f32, %x2: f32, %x3: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) + +template.decl @llm.ssm_conv.prefill_finish(%state: buffer, %filter: buffer, %edges: buffer, %output: buffer, %cache: buffer) + +config.decl @llm.ssm_conv.generic.d_inner : %value: index where [range(%value, 32, 65536), mul(%value, 32)] + +template.def<@llm.ssm_conv.prefill_finish> device @llm_ssm_conv_prefill_finish(%state: buffer, %filter: buffer, %edges: buffer, %output: buffer, %cache: buffer) { + %width = config.get @llm.ssm_conv.generic.d_inner : index + %zero = index.constant 0 : index + %one = index.constant 1 : index + %two = index.constant 2 : index + %three = index.constant 3 : index + %four = index.constant 4 : index + %five = index.constant 5 : index + %threads = index.constant 256 : index + %base = index.constant 0 : offset + %group = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %channel0 = index.madd %group, %threads, %lane : index + %valid = index.cmp ult, %channel0, %width : index + scf.if %valid { + %channel, %bounded_width = index.assume %channel0, %width [lt(%channel0, %width)] : index, index + %state_view = buffer.view %state[%base] : buffer -> view<[%bounded_width]x3xf32> + %filter_view = buffer.view %filter[%base] : buffer -> view<[%bounded_width]x4xf32> + %edge_view = buffer.view %edges[%base] : buffer -> view<[%bounded_width]x6xf32> + %output_view = buffer.view %output[%base] : buffer -> view<512x[%bounded_width]xf32> + %cache_view = buffer.view %cache[%base] : buffer -> view<[%bounded_width]x3xf32> + %s0 = view.load %state_view[%channel, %zero] : view<[%bounded_width]x3xf32> -> f32 + %s1 = view.load %state_view[%channel, %one] : view<[%bounded_width]x3xf32> -> f32 + %s2 = view.load %state_view[%channel, %two] : view<[%bounded_width]x3xf32> -> f32 + %x0 = view.load %edge_view[%channel, %zero] : view<[%bounded_width]x6xf32> -> f32 + %x1 = view.load %edge_view[%channel, %one] : view<[%bounded_width]x6xf32> -> f32 + %x2 = view.load %edge_view[%channel, %two] : view<[%bounded_width]x6xf32> -> f32 + %w0 = view.load %filter_view[%channel, %zero] : view<[%bounded_width]x4xf32> -> f32 + %w1 = view.load %filter_view[%channel, %one] : view<[%bounded_width]x4xf32> -> f32 + %w2 = view.load %filter_view[%channel, %two] : view<[%bounded_width]x4xf32> -> f32 + %w3 = view.load %filter_view[%channel, %three] : view<[%bounded_width]x4xf32> -> f32 + %y0 = template.apply<@llm.ssm_conv.dconv4_silu>(%s0, %s1, %s2, %x0, %w0, %w1, %w2, %w3) : (f32, f32, f32, f32, f32, f32, f32, f32) -> (f32) + %y1 = template.apply<@llm.ssm_conv.dconv4_silu>(%s1, %s2, %x0, %x1, %w0, %w1, %w2, %w3) : (f32, f32, f32, f32, f32, f32, f32, f32) -> (f32) + %y2 = template.apply<@llm.ssm_conv.dconv4_silu>(%s2, %x0, %x1, %x2, %w0, %w1, %w2, %w3) : (f32, f32, f32, f32, f32, f32, f32, f32) -> (f32) + view.store %y0, %output_view[%zero, %channel] : f32, view<512x[%bounded_width]xf32> + view.store %y1, %output_view[%one, %channel] : f32, view<512x[%bounded_width]xf32> + view.store %y2, %output_view[%two, %channel] : f32, view<512x[%bounded_width]xf32> + %tail0 = view.load %edge_view[%channel, %three] : view<[%bounded_width]x6xf32> -> f32 + %tail1 = view.load %edge_view[%channel, %four] : view<[%bounded_width]x6xf32> -> f32 + %tail2 = view.load %edge_view[%channel, %five] : view<[%bounded_width]x6xf32> -> f32 + view.store %tail0, %cache_view[%channel, %zero] : f32, view<[%bounded_width]x3xf32> + view.store %tail1, %cache_view[%channel, %one] : f32, view<[%bounded_width]x3xf32> + view.store %tail2, %cache_view[%channel, %two] : f32, view<[%bounded_width]x3xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/prefill_window.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/prefill_window.loom new file mode 100644 index 000000000000..a272d82db7ac --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/prefill_window.loom @@ -0,0 +1,708 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.prefill_window(%packet_width: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer) + +config.decl @llm.ssm_conv.prefill.d_conv : %value: index where [range(%value, 4, 4)] + +config.decl @llm.ssm_conv.prefill.d_inner : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +config.decl @llm.ssm_conv.prefill.n_t : %value: index where [range(%value, 512, 512)] + +config.decl @llm.ssm_conv.prefill.n_s : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.prefill.state_row_stride : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +config.decl @llm.ssm_conv.prefill.x_row_stride : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +template.def<@llm.ssm_conv.prefill_window> device @llm_ssm_conv_prefill_window(%packet_width: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c5 = index.constant 5 : index + %c6 = index.constant 6 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c9 = index.constant 9 : index + %c10 = index.constant 10 : index + %c11 = index.constant 11 : index + %c12 = index.constant 12 : index + %c13 = index.constant 13 : index + %c14 = index.constant 14 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c17 = index.constant 17 : index + %c18 = index.constant 18 : index + %c19 = index.constant 19 : 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 + %c26 = index.constant 26 : index + %c27 = index.constant 27 : index + %c28 = index.constant 28 : index + %c29 = index.constant 29 : index + %c30 = index.constant 30 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c61 = index.constant 61 : index + %channel_limit = index.constant 10240 : index + %filter_limit = index.constant 163840 : index + %flat_limit = index.constant 1073741824 : index + %channel_last = index.sub %channel_limit, %packet_width : index + %packet_bytes = index.mul %packet_width, %c16 : index + %filter_last = index.sub %filter_limit, %packet_bytes : index + %flat_last = index.sub %flat_limit, %packet_width : index + %c0_f32 = vector.constant 0.0 : vector<[%packet_width]xf32> + + %d_inner0 = config.get @llm.ssm_conv.prefill.d_inner : index + %n_t0 = config.get @llm.ssm_conv.prefill.n_t : index + %n_s0 = config.get @llm.ssm_conv.prefill.n_s : index + %state_row_stride0 = config.get @llm.ssm_conv.prefill.state_row_stride : index + %x_row_stride0 = config.get @llm.ssm_conv.prefill.x_row_stride : index + %d_inner = index.assume %d_inner0 [range(%d_inner0, 8192, 10240), mul(%d_inner0, 32)] : index + %n_t = index.assume %n_t0 [range(%n_t0, 512, 512)] : index + %n_s = index.assume %n_s0 [range(%n_s0, 1, 1)] : index + %state_row_stride = index.assume %state_row_stride0 [range(%state_row_stride0, 8192, 10240), mul(%state_row_stride0, 32)] : index + %x_row_stride = index.assume %x_row_stride0 [range(%x_row_stride0, 8192, 10240), mul(%x_row_stride0, 32)] : index + + %group_x0 = kernel.workgroup.id : index + %sequence0 = kernel.workgroup.id : index + %local0 = kernel.workitem.id : index + %group_x = index.assume %group_x0 [range(%group_x0, 0, 319)] : index + %sequence = index.assume %sequence0 [range(%sequence0, 0, 0)] : index + %local = index.assume %local0 [range(%local0, 0, 1023)] : index + %wave0 = index.div %local, %c32 : index + %wave = index.assume %wave0 [range(%wave0, 0, 31)] : index + %lane0 = index.rem %local, %c32 : index + %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index + %first_wave = index.cmp eq, %wave, %c0 : index + + %group_channels = index.mul %packet_width, %c32 : index + %channel_base0 = index.mul %group_x, %group_channels : index + %channel_base = index.assume %channel_base0 [range(%channel_base0, 0, 10208), mul(%channel_base0, 32)] : index + %lane_channel = index.mul %lane, %packet_width : index + %channel0 = index.add %channel_base, %lane_channel : index + %channel = index.assume %channel0 [range(%channel0, 0, %channel_last), mul(%channel0, %packet_width)] : index + %wave_token_base0 = index.mul %wave, %c16 : index + %wave_token_base = index.assume %wave_token_base0 [range(%wave_token_base0, 0, 496), mul(%wave_token_base0, 16)] : index + %prior_wave = scf.if %first_wave -> (index) { + scf.yield %c0 : index + } else { + %prior_wave0 = index.sub %wave, %c1 : index + %prior_wave1 = index.assume %prior_wave0 [range(%prior_wave0, 0, 30)] : index + scf.yield %prior_wave1 : index + } + %prior_wave_base0 = index.mul %prior_wave, %c16 : index + %prior_wave_base = index.assume %prior_wave_base0 [range(%prior_wave_base0, 0, 480), mul(%prior_wave_base0, 16)] : index + + %state_g = buffer.assume.memory_space %state : buffer + %x_g = buffer.assume.memory_space %x : buffer + %filt_g = buffer.assume.memory_space %filt : buffer + %dst_g = buffer.assume.memory_space %dst : buffer + // dst deliberately does not participate: its physical span may equal x. + %state_na, %x_na, %filt_na = buffer.assume.noalias %state_g, %x_g, %filt_g : buffer, buffer, buffer + %state_v = buffer.view %state_na[%base] : buffer -> view<1073741824xf32> + %x_v = buffer.view %x_na[%base] : buffer -> view<1073741824xf32> + %dst_v = buffer.view %dst_g[%base] : buffer -> view<1073741824xf32> + + %x_sequence_span0 = index.mul %n_t, %x_row_stride : index + %x_sequence_span = index.assume %x_sequence_span0 [range(%x_sequence_span0, 4194304, 5242880), mul(%x_sequence_span0, 16384)] : index + %x_sequence_base0 = index.mul %sequence, %x_sequence_span : index + %x_sequence_base = index.assume %x_sequence_base0 [range(%x_sequence_base0, 0, %flat_last)] : index + + %state_sequence_rows0 = index.add %n_t, %c3 : index + %state_sequence_rows = index.assume %state_sequence_rows0 [range(%state_sequence_rows0, 515, 515)] : index + %state_sequence_span0 = index.mul %state_sequence_rows, %state_row_stride : index + %state_sequence_span = index.assume %state_sequence_span0 [range(%state_sequence_span0, 4218880, 5273600), mul(%state_sequence_span0, 16480)] : index + %state_sequence_base0 = index.mul %sequence, %state_sequence_span : index + %state_sequence_base = index.assume %state_sequence_base0 [range(%state_sequence_base0, 0, %flat_last)] : index + + // Snapshot slot k maps to x row wave_base-3+k; wave zero gets its first three taps from carried state. + // Other waves preload all 19 slots, and rows 0..60 use the immutable CONCAT window. + %xr0, %xr1, %xr2 = scf.if %first_wave -> (vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>) { + scf.yield %c0_f32, %c0_f32, %c0_f32 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } else { + %x_row_00 = index.add %prior_wave_base, %c13 : index + %x_row_0 = index.assume %x_row_00 [range(%x_row_00, 13, 508)] : index + %x_row_10 = index.add %prior_wave_base, %c14 : index + %x_row_1 = index.assume %x_row_10 [range(%x_row_10, 14, 509)] : index + %x_row_20 = index.add %prior_wave_base, %c15 : index + %x_row_2 = index.assume %x_row_20 [range(%x_row_20, 15, 510)] : index + %early_is_snapshotted = index.cmp ult, %wave, %c4 : index + %early_value_0, %early_value_1, %early_value_2 = scf.if %early_is_snapshotted -> (vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>) { + %snapshot_x_row_early_0 = index.assume %x_row_0 [range(%x_row_0, 0, 60)] : index + %snapshot_row_early_00 = index.add %snapshot_x_row_early_0, %c3 : index + %snapshot_row_early_0 = index.assume %snapshot_row_early_00 [range(%snapshot_row_early_00, 3, 63)] : index + %snapshot_row_early_0_offset0 = index.mul %snapshot_row_early_0, %state_row_stride : index + %snapshot_row_early_0_offset1 = index.add %state_sequence_base, %snapshot_row_early_0_offset0 : index + %snapshot_index_early_00 = index.add %snapshot_row_early_0_offset1, %channel : index + %snapshot_index_early_0 = index.assume %snapshot_index_early_00 [range(%snapshot_index_early_00, 0, %flat_last)] : index + %early_snapshot_value_0 = vector.load %state_v[%snapshot_index_early_0] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_early_1 = index.assume %x_row_1 [range(%x_row_1, 0, 60)] : index + %snapshot_row_early_10 = index.add %snapshot_x_row_early_1, %c3 : index + %snapshot_row_early_1 = index.assume %snapshot_row_early_10 [range(%snapshot_row_early_10, 3, 63)] : index + %snapshot_row_early_1_offset0 = index.mul %snapshot_row_early_1, %state_row_stride : index + %snapshot_row_early_1_offset1 = index.add %state_sequence_base, %snapshot_row_early_1_offset0 : index + %snapshot_index_early_10 = index.add %snapshot_row_early_1_offset1, %channel : index + %snapshot_index_early_1 = index.assume %snapshot_index_early_10 [range(%snapshot_index_early_10, 0, %flat_last)] : index + %early_snapshot_value_1 = vector.load %state_v[%snapshot_index_early_1] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_early_2 = index.assume %x_row_2 [range(%x_row_2, 0, 60)] : index + %snapshot_row_early_20 = index.add %snapshot_x_row_early_2, %c3 : index + %snapshot_row_early_2 = index.assume %snapshot_row_early_20 [range(%snapshot_row_early_20, 3, 63)] : index + %snapshot_row_early_2_offset0 = index.mul %snapshot_row_early_2, %state_row_stride : index + %snapshot_row_early_2_offset1 = index.add %state_sequence_base, %snapshot_row_early_2_offset0 : index + %snapshot_index_early_20 = index.add %snapshot_row_early_2_offset1, %channel : index + %snapshot_index_early_2 = index.assume %snapshot_index_early_20 [range(%snapshot_index_early_20, 0, %flat_last)] : index + %early_snapshot_value_2 = vector.load %state_v[%snapshot_index_early_2] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %early_snapshot_value_0, %early_snapshot_value_1, %early_snapshot_value_2 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } else { + %x_row_early_00 = index.mul %x_row_0, %x_row_stride : index + %x_row_early_01 = index.add %x_sequence_base, %x_row_early_00 : index + %x_index_early_00 = index.add %x_row_early_01, %channel : index + %x_index_early_0 = index.assume %x_index_early_00 [range(%x_index_early_00, 0, %flat_last)] : index + %early_raw_value_0 = vector.load %x_v[%x_index_early_0] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_early_10 = index.mul %x_row_1, %x_row_stride : index + %x_row_early_11 = index.add %x_sequence_base, %x_row_early_10 : index + %x_index_early_10 = index.add %x_row_early_11, %channel : index + %x_index_early_1 = index.assume %x_index_early_10 [range(%x_index_early_10, 0, %flat_last)] : index + %early_raw_value_1 = vector.load %x_v[%x_index_early_1] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_early_20 = index.mul %x_row_2, %x_row_stride : index + %x_row_early_21 = index.add %x_sequence_base, %x_row_early_20 : index + %x_index_early_20 = index.add %x_row_early_21, %channel : index + %x_index_early_2 = index.assume %x_index_early_20 [range(%x_index_early_20, 0, %flat_last)] : index + %early_raw_value_2 = vector.load %x_v[%x_index_early_2] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %early_raw_value_0, %early_raw_value_1, %early_raw_value_2 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } + scf.yield %early_value_0, %early_value_1, %early_value_2 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } + %x_row_common_30 = index.add %wave_token_base, %c0 : index + %x_row_common_3 = index.assume %x_row_common_30 [range(%x_row_common_30, 0, 496)] : index + %x_row_common_40 = index.add %wave_token_base, %c1 : index + %x_row_common_4 = index.assume %x_row_common_40 [range(%x_row_common_40, 1, 497)] : index + %x_row_common_50 = index.add %wave_token_base, %c2 : index + %x_row_common_5 = index.assume %x_row_common_50 [range(%x_row_common_50, 2, 498)] : index + %x_row_common_60 = index.add %wave_token_base, %c3 : index + %x_row_common_6 = index.assume %x_row_common_60 [range(%x_row_common_60, 3, 499)] : index + %x_row_common_70 = index.add %wave_token_base, %c4 : index + %x_row_common_7 = index.assume %x_row_common_70 [range(%x_row_common_70, 4, 500)] : index + %x_row_common_80 = index.add %wave_token_base, %c5 : index + %x_row_common_8 = index.assume %x_row_common_80 [range(%x_row_common_80, 5, 501)] : index + %x_row_common_90 = index.add %wave_token_base, %c6 : index + %x_row_common_9 = index.assume %x_row_common_90 [range(%x_row_common_90, 6, 502)] : index + %x_row_common_100 = index.add %wave_token_base, %c7 : index + %x_row_common_10 = index.assume %x_row_common_100 [range(%x_row_common_100, 7, 503)] : index + %x_row_common_110 = index.add %wave_token_base, %c8 : index + %x_row_common_11 = index.assume %x_row_common_110 [range(%x_row_common_110, 8, 504)] : index + %x_row_common_120 = index.add %wave_token_base, %c9 : index + %x_row_common_12 = index.assume %x_row_common_120 [range(%x_row_common_120, 9, 505)] : index + %x_row_common_130 = index.add %wave_token_base, %c10 : index + %x_row_common_13 = index.assume %x_row_common_130 [range(%x_row_common_130, 10, 506)] : index + %x_row_common_140 = index.add %wave_token_base, %c11 : index + %x_row_common_14 = index.assume %x_row_common_140 [range(%x_row_common_140, 11, 507)] : index + %x_row_common_150 = index.add %wave_token_base, %c12 : index + %x_row_common_15 = index.assume %x_row_common_150 [range(%x_row_common_150, 12, 508)] : index + %common_low_is_snapshotted = index.cmp ult, %wave, %c4 : index + %xr3, %xr4, %xr5, %xr6, %xr7, %xr8, %xr9, %xr10, %xr11, %xr12, %xr13, %xr14, %xr15 = scf.if %common_low_is_snapshotted -> (vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>) { + %snapshot_x_row_common_low_3 = index.assume %x_row_common_3 [range(%x_row_common_3, 0, 60)] : index + %snapshot_row_common_low_30 = index.add %snapshot_x_row_common_low_3, %c3 : index + %snapshot_row_common_low_3 = index.assume %snapshot_row_common_low_30 [range(%snapshot_row_common_low_30, 3, 63)] : index + %snapshot_row_common_low_3_offset0 = index.mul %snapshot_row_common_low_3, %state_row_stride : index + %snapshot_row_common_low_3_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_3_offset0 : index + %snapshot_index_common_low_30 = index.add %snapshot_row_common_low_3_offset1, %channel : index + %snapshot_index_common_low_3 = index.assume %snapshot_index_common_low_30 [range(%snapshot_index_common_low_30, 0, %flat_last)] : index + %common_low_snapshot_value_3 = vector.load %state_v[%snapshot_index_common_low_3] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_4 = index.assume %x_row_common_4 [range(%x_row_common_4, 0, 60)] : index + %snapshot_row_common_low_40 = index.add %snapshot_x_row_common_low_4, %c3 : index + %snapshot_row_common_low_4 = index.assume %snapshot_row_common_low_40 [range(%snapshot_row_common_low_40, 3, 63)] : index + %snapshot_row_common_low_4_offset0 = index.mul %snapshot_row_common_low_4, %state_row_stride : index + %snapshot_row_common_low_4_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_4_offset0 : index + %snapshot_index_common_low_40 = index.add %snapshot_row_common_low_4_offset1, %channel : index + %snapshot_index_common_low_4 = index.assume %snapshot_index_common_low_40 [range(%snapshot_index_common_low_40, 0, %flat_last)] : index + %common_low_snapshot_value_4 = vector.load %state_v[%snapshot_index_common_low_4] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_5 = index.assume %x_row_common_5 [range(%x_row_common_5, 0, 60)] : index + %snapshot_row_common_low_50 = index.add %snapshot_x_row_common_low_5, %c3 : index + %snapshot_row_common_low_5 = index.assume %snapshot_row_common_low_50 [range(%snapshot_row_common_low_50, 3, 63)] : index + %snapshot_row_common_low_5_offset0 = index.mul %snapshot_row_common_low_5, %state_row_stride : index + %snapshot_row_common_low_5_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_5_offset0 : index + %snapshot_index_common_low_50 = index.add %snapshot_row_common_low_5_offset1, %channel : index + %snapshot_index_common_low_5 = index.assume %snapshot_index_common_low_50 [range(%snapshot_index_common_low_50, 0, %flat_last)] : index + %common_low_snapshot_value_5 = vector.load %state_v[%snapshot_index_common_low_5] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_6 = index.assume %x_row_common_6 [range(%x_row_common_6, 0, 60)] : index + %snapshot_row_common_low_60 = index.add %snapshot_x_row_common_low_6, %c3 : index + %snapshot_row_common_low_6 = index.assume %snapshot_row_common_low_60 [range(%snapshot_row_common_low_60, 3, 63)] : index + %snapshot_row_common_low_6_offset0 = index.mul %snapshot_row_common_low_6, %state_row_stride : index + %snapshot_row_common_low_6_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_6_offset0 : index + %snapshot_index_common_low_60 = index.add %snapshot_row_common_low_6_offset1, %channel : index + %snapshot_index_common_low_6 = index.assume %snapshot_index_common_low_60 [range(%snapshot_index_common_low_60, 0, %flat_last)] : index + %common_low_snapshot_value_6 = vector.load %state_v[%snapshot_index_common_low_6] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_7 = index.assume %x_row_common_7 [range(%x_row_common_7, 0, 60)] : index + %snapshot_row_common_low_70 = index.add %snapshot_x_row_common_low_7, %c3 : index + %snapshot_row_common_low_7 = index.assume %snapshot_row_common_low_70 [range(%snapshot_row_common_low_70, 3, 63)] : index + %snapshot_row_common_low_7_offset0 = index.mul %snapshot_row_common_low_7, %state_row_stride : index + %snapshot_row_common_low_7_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_7_offset0 : index + %snapshot_index_common_low_70 = index.add %snapshot_row_common_low_7_offset1, %channel : index + %snapshot_index_common_low_7 = index.assume %snapshot_index_common_low_70 [range(%snapshot_index_common_low_70, 0, %flat_last)] : index + %common_low_snapshot_value_7 = vector.load %state_v[%snapshot_index_common_low_7] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_8 = index.assume %x_row_common_8 [range(%x_row_common_8, 0, 60)] : index + %snapshot_row_common_low_80 = index.add %snapshot_x_row_common_low_8, %c3 : index + %snapshot_row_common_low_8 = index.assume %snapshot_row_common_low_80 [range(%snapshot_row_common_low_80, 3, 63)] : index + %snapshot_row_common_low_8_offset0 = index.mul %snapshot_row_common_low_8, %state_row_stride : index + %snapshot_row_common_low_8_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_8_offset0 : index + %snapshot_index_common_low_80 = index.add %snapshot_row_common_low_8_offset1, %channel : index + %snapshot_index_common_low_8 = index.assume %snapshot_index_common_low_80 [range(%snapshot_index_common_low_80, 0, %flat_last)] : index + %common_low_snapshot_value_8 = vector.load %state_v[%snapshot_index_common_low_8] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_9 = index.assume %x_row_common_9 [range(%x_row_common_9, 0, 60)] : index + %snapshot_row_common_low_90 = index.add %snapshot_x_row_common_low_9, %c3 : index + %snapshot_row_common_low_9 = index.assume %snapshot_row_common_low_90 [range(%snapshot_row_common_low_90, 3, 63)] : index + %snapshot_row_common_low_9_offset0 = index.mul %snapshot_row_common_low_9, %state_row_stride : index + %snapshot_row_common_low_9_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_9_offset0 : index + %snapshot_index_common_low_90 = index.add %snapshot_row_common_low_9_offset1, %channel : index + %snapshot_index_common_low_9 = index.assume %snapshot_index_common_low_90 [range(%snapshot_index_common_low_90, 0, %flat_last)] : index + %common_low_snapshot_value_9 = vector.load %state_v[%snapshot_index_common_low_9] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_10 = index.assume %x_row_common_10 [range(%x_row_common_10, 0, 60)] : index + %snapshot_row_common_low_100 = index.add %snapshot_x_row_common_low_10, %c3 : index + %snapshot_row_common_low_10 = index.assume %snapshot_row_common_low_100 [range(%snapshot_row_common_low_100, 3, 63)] : index + %snapshot_row_common_low_10_offset0 = index.mul %snapshot_row_common_low_10, %state_row_stride : index + %snapshot_row_common_low_10_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_10_offset0 : index + %snapshot_index_common_low_100 = index.add %snapshot_row_common_low_10_offset1, %channel : index + %snapshot_index_common_low_10 = index.assume %snapshot_index_common_low_100 [range(%snapshot_index_common_low_100, 0, %flat_last)] : index + %common_low_snapshot_value_10 = vector.load %state_v[%snapshot_index_common_low_10] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_11 = index.assume %x_row_common_11 [range(%x_row_common_11, 0, 60)] : index + %snapshot_row_common_low_110 = index.add %snapshot_x_row_common_low_11, %c3 : index + %snapshot_row_common_low_11 = index.assume %snapshot_row_common_low_110 [range(%snapshot_row_common_low_110, 3, 63)] : index + %snapshot_row_common_low_11_offset0 = index.mul %snapshot_row_common_low_11, %state_row_stride : index + %snapshot_row_common_low_11_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_11_offset0 : index + %snapshot_index_common_low_110 = index.add %snapshot_row_common_low_11_offset1, %channel : index + %snapshot_index_common_low_11 = index.assume %snapshot_index_common_low_110 [range(%snapshot_index_common_low_110, 0, %flat_last)] : index + %common_low_snapshot_value_11 = vector.load %state_v[%snapshot_index_common_low_11] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_12 = index.assume %x_row_common_12 [range(%x_row_common_12, 0, 60)] : index + %snapshot_row_common_low_120 = index.add %snapshot_x_row_common_low_12, %c3 : index + %snapshot_row_common_low_12 = index.assume %snapshot_row_common_low_120 [range(%snapshot_row_common_low_120, 3, 63)] : index + %snapshot_row_common_low_12_offset0 = index.mul %snapshot_row_common_low_12, %state_row_stride : index + %snapshot_row_common_low_12_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_12_offset0 : index + %snapshot_index_common_low_120 = index.add %snapshot_row_common_low_12_offset1, %channel : index + %snapshot_index_common_low_12 = index.assume %snapshot_index_common_low_120 [range(%snapshot_index_common_low_120, 0, %flat_last)] : index + %common_low_snapshot_value_12 = vector.load %state_v[%snapshot_index_common_low_12] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_13 = index.assume %x_row_common_13 [range(%x_row_common_13, 0, 60)] : index + %snapshot_row_common_low_130 = index.add %snapshot_x_row_common_low_13, %c3 : index + %snapshot_row_common_low_13 = index.assume %snapshot_row_common_low_130 [range(%snapshot_row_common_low_130, 3, 63)] : index + %snapshot_row_common_low_13_offset0 = index.mul %snapshot_row_common_low_13, %state_row_stride : index + %snapshot_row_common_low_13_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_13_offset0 : index + %snapshot_index_common_low_130 = index.add %snapshot_row_common_low_13_offset1, %channel : index + %snapshot_index_common_low_13 = index.assume %snapshot_index_common_low_130 [range(%snapshot_index_common_low_130, 0, %flat_last)] : index + %common_low_snapshot_value_13 = vector.load %state_v[%snapshot_index_common_low_13] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_14 = index.assume %x_row_common_14 [range(%x_row_common_14, 0, 60)] : index + %snapshot_row_common_low_140 = index.add %snapshot_x_row_common_low_14, %c3 : index + %snapshot_row_common_low_14 = index.assume %snapshot_row_common_low_140 [range(%snapshot_row_common_low_140, 3, 63)] : index + %snapshot_row_common_low_14_offset0 = index.mul %snapshot_row_common_low_14, %state_row_stride : index + %snapshot_row_common_low_14_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_14_offset0 : index + %snapshot_index_common_low_140 = index.add %snapshot_row_common_low_14_offset1, %channel : index + %snapshot_index_common_low_14 = index.assume %snapshot_index_common_low_140 [range(%snapshot_index_common_low_140, 0, %flat_last)] : index + %common_low_snapshot_value_14 = vector.load %state_v[%snapshot_index_common_low_14] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_low_15 = index.assume %x_row_common_15 [range(%x_row_common_15, 0, 60)] : index + %snapshot_row_common_low_150 = index.add %snapshot_x_row_common_low_15, %c3 : index + %snapshot_row_common_low_15 = index.assume %snapshot_row_common_low_150 [range(%snapshot_row_common_low_150, 3, 63)] : index + %snapshot_row_common_low_15_offset0 = index.mul %snapshot_row_common_low_15, %state_row_stride : index + %snapshot_row_common_low_15_offset1 = index.add %state_sequence_base, %snapshot_row_common_low_15_offset0 : index + %snapshot_index_common_low_150 = index.add %snapshot_row_common_low_15_offset1, %channel : index + %snapshot_index_common_low_15 = index.assume %snapshot_index_common_low_150 [range(%snapshot_index_common_low_150, 0, %flat_last)] : index + %common_low_snapshot_value_15 = vector.load %state_v[%snapshot_index_common_low_15] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %common_low_snapshot_value_3, %common_low_snapshot_value_4, %common_low_snapshot_value_5, %common_low_snapshot_value_6, %common_low_snapshot_value_7, %common_low_snapshot_value_8, %common_low_snapshot_value_9, %common_low_snapshot_value_10, %common_low_snapshot_value_11, %common_low_snapshot_value_12, %common_low_snapshot_value_13, %common_low_snapshot_value_14, %common_low_snapshot_value_15 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } else { + %x_row_common_low_30 = index.mul %x_row_common_3, %x_row_stride : index + %x_row_common_low_31 = index.add %x_sequence_base, %x_row_common_low_30 : index + %x_index_common_low_30 = index.add %x_row_common_low_31, %channel : index + %x_index_common_low_3 = index.assume %x_index_common_low_30 [range(%x_index_common_low_30, 0, %flat_last)] : index + %common_low_raw_value_3 = vector.load %x_v[%x_index_common_low_3] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_40 = index.mul %x_row_common_4, %x_row_stride : index + %x_row_common_low_41 = index.add %x_sequence_base, %x_row_common_low_40 : index + %x_index_common_low_40 = index.add %x_row_common_low_41, %channel : index + %x_index_common_low_4 = index.assume %x_index_common_low_40 [range(%x_index_common_low_40, 0, %flat_last)] : index + %common_low_raw_value_4 = vector.load %x_v[%x_index_common_low_4] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_50 = index.mul %x_row_common_5, %x_row_stride : index + %x_row_common_low_51 = index.add %x_sequence_base, %x_row_common_low_50 : index + %x_index_common_low_50 = index.add %x_row_common_low_51, %channel : index + %x_index_common_low_5 = index.assume %x_index_common_low_50 [range(%x_index_common_low_50, 0, %flat_last)] : index + %common_low_raw_value_5 = vector.load %x_v[%x_index_common_low_5] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_60 = index.mul %x_row_common_6, %x_row_stride : index + %x_row_common_low_61 = index.add %x_sequence_base, %x_row_common_low_60 : index + %x_index_common_low_60 = index.add %x_row_common_low_61, %channel : index + %x_index_common_low_6 = index.assume %x_index_common_low_60 [range(%x_index_common_low_60, 0, %flat_last)] : index + %common_low_raw_value_6 = vector.load %x_v[%x_index_common_low_6] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_70 = index.mul %x_row_common_7, %x_row_stride : index + %x_row_common_low_71 = index.add %x_sequence_base, %x_row_common_low_70 : index + %x_index_common_low_70 = index.add %x_row_common_low_71, %channel : index + %x_index_common_low_7 = index.assume %x_index_common_low_70 [range(%x_index_common_low_70, 0, %flat_last)] : index + %common_low_raw_value_7 = vector.load %x_v[%x_index_common_low_7] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_80 = index.mul %x_row_common_8, %x_row_stride : index + %x_row_common_low_81 = index.add %x_sequence_base, %x_row_common_low_80 : index + %x_index_common_low_80 = index.add %x_row_common_low_81, %channel : index + %x_index_common_low_8 = index.assume %x_index_common_low_80 [range(%x_index_common_low_80, 0, %flat_last)] : index + %common_low_raw_value_8 = vector.load %x_v[%x_index_common_low_8] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_90 = index.mul %x_row_common_9, %x_row_stride : index + %x_row_common_low_91 = index.add %x_sequence_base, %x_row_common_low_90 : index + %x_index_common_low_90 = index.add %x_row_common_low_91, %channel : index + %x_index_common_low_9 = index.assume %x_index_common_low_90 [range(%x_index_common_low_90, 0, %flat_last)] : index + %common_low_raw_value_9 = vector.load %x_v[%x_index_common_low_9] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_100 = index.mul %x_row_common_10, %x_row_stride : index + %x_row_common_low_101 = index.add %x_sequence_base, %x_row_common_low_100 : index + %x_index_common_low_100 = index.add %x_row_common_low_101, %channel : index + %x_index_common_low_10 = index.assume %x_index_common_low_100 [range(%x_index_common_low_100, 0, %flat_last)] : index + %common_low_raw_value_10 = vector.load %x_v[%x_index_common_low_10] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_110 = index.mul %x_row_common_11, %x_row_stride : index + %x_row_common_low_111 = index.add %x_sequence_base, %x_row_common_low_110 : index + %x_index_common_low_110 = index.add %x_row_common_low_111, %channel : index + %x_index_common_low_11 = index.assume %x_index_common_low_110 [range(%x_index_common_low_110, 0, %flat_last)] : index + %common_low_raw_value_11 = vector.load %x_v[%x_index_common_low_11] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_120 = index.mul %x_row_common_12, %x_row_stride : index + %x_row_common_low_121 = index.add %x_sequence_base, %x_row_common_low_120 : index + %x_index_common_low_120 = index.add %x_row_common_low_121, %channel : index + %x_index_common_low_12 = index.assume %x_index_common_low_120 [range(%x_index_common_low_120, 0, %flat_last)] : index + %common_low_raw_value_12 = vector.load %x_v[%x_index_common_low_12] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_130 = index.mul %x_row_common_13, %x_row_stride : index + %x_row_common_low_131 = index.add %x_sequence_base, %x_row_common_low_130 : index + %x_index_common_low_130 = index.add %x_row_common_low_131, %channel : index + %x_index_common_low_13 = index.assume %x_index_common_low_130 [range(%x_index_common_low_130, 0, %flat_last)] : index + %common_low_raw_value_13 = vector.load %x_v[%x_index_common_low_13] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_140 = index.mul %x_row_common_14, %x_row_stride : index + %x_row_common_low_141 = index.add %x_sequence_base, %x_row_common_low_140 : index + %x_index_common_low_140 = index.add %x_row_common_low_141, %channel : index + %x_index_common_low_14 = index.assume %x_index_common_low_140 [range(%x_index_common_low_140, 0, %flat_last)] : index + %common_low_raw_value_14 = vector.load %x_v[%x_index_common_low_14] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_low_150 = index.mul %x_row_common_15, %x_row_stride : index + %x_row_common_low_151 = index.add %x_sequence_base, %x_row_common_low_150 : index + %x_index_common_low_150 = index.add %x_row_common_low_151, %channel : index + %x_index_common_low_15 = index.assume %x_index_common_low_150 [range(%x_index_common_low_150, 0, %flat_last)] : index + %common_low_raw_value_15 = vector.load %x_v[%x_index_common_low_15] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %common_low_raw_value_3, %common_low_raw_value_4, %common_low_raw_value_5, %common_low_raw_value_6, %common_low_raw_value_7, %common_low_raw_value_8, %common_low_raw_value_9, %common_low_raw_value_10, %common_low_raw_value_11, %common_low_raw_value_12, %common_low_raw_value_13, %common_low_raw_value_14, %common_low_raw_value_15 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } + %x_row_common_160 = index.add %wave_token_base, %c13 : index + %x_row_common_16 = index.assume %x_row_common_160 [range(%x_row_common_160, 13, 509)] : index + %x_row_common_170 = index.add %wave_token_base, %c14 : index + %x_row_common_17 = index.assume %x_row_common_170 [range(%x_row_common_170, 14, 510)] : index + %x_row_common_180 = index.add %wave_token_base, %c15 : index + %x_row_common_18 = index.assume %x_row_common_180 [range(%x_row_common_180, 15, 511)] : index + %common_high_is_snapshotted = index.cmp ult, %wave, %c3 : index + %xr16, %xr17, %xr18 = scf.if %common_high_is_snapshotted -> (vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32>) { + %snapshot_x_row_common_high_16 = index.assume %x_row_common_16 [range(%x_row_common_16, 0, 60)] : index + %snapshot_row_common_high_160 = index.add %snapshot_x_row_common_high_16, %c3 : index + %snapshot_row_common_high_16 = index.assume %snapshot_row_common_high_160 [range(%snapshot_row_common_high_160, 3, 63)] : index + %snapshot_row_common_high_16_offset0 = index.mul %snapshot_row_common_high_16, %state_row_stride : index + %snapshot_row_common_high_16_offset1 = index.add %state_sequence_base, %snapshot_row_common_high_16_offset0 : index + %snapshot_index_common_high_160 = index.add %snapshot_row_common_high_16_offset1, %channel : index + %snapshot_index_common_high_16 = index.assume %snapshot_index_common_high_160 [range(%snapshot_index_common_high_160, 0, %flat_last)] : index + %common_high_snapshot_value_16 = vector.load %state_v[%snapshot_index_common_high_16] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_high_17 = index.assume %x_row_common_17 [range(%x_row_common_17, 0, 60)] : index + %snapshot_row_common_high_170 = index.add %snapshot_x_row_common_high_17, %c3 : index + %snapshot_row_common_high_17 = index.assume %snapshot_row_common_high_170 [range(%snapshot_row_common_high_170, 3, 63)] : index + %snapshot_row_common_high_17_offset0 = index.mul %snapshot_row_common_high_17, %state_row_stride : index + %snapshot_row_common_high_17_offset1 = index.add %state_sequence_base, %snapshot_row_common_high_17_offset0 : index + %snapshot_index_common_high_170 = index.add %snapshot_row_common_high_17_offset1, %channel : index + %snapshot_index_common_high_17 = index.assume %snapshot_index_common_high_170 [range(%snapshot_index_common_high_170, 0, %flat_last)] : index + %common_high_snapshot_value_17 = vector.load %state_v[%snapshot_index_common_high_17] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %snapshot_x_row_common_high_18 = index.assume %x_row_common_18 [range(%x_row_common_18, 0, 60)] : index + %snapshot_row_common_high_180 = index.add %snapshot_x_row_common_high_18, %c3 : index + %snapshot_row_common_high_18 = index.assume %snapshot_row_common_high_180 [range(%snapshot_row_common_high_180, 3, 63)] : index + %snapshot_row_common_high_18_offset0 = index.mul %snapshot_row_common_high_18, %state_row_stride : index + %snapshot_row_common_high_18_offset1 = index.add %state_sequence_base, %snapshot_row_common_high_18_offset0 : index + %snapshot_index_common_high_180 = index.add %snapshot_row_common_high_18_offset1, %channel : index + %snapshot_index_common_high_18 = index.assume %snapshot_index_common_high_180 [range(%snapshot_index_common_high_180, 0, %flat_last)] : index + %common_high_snapshot_value_18 = vector.load %state_v[%snapshot_index_common_high_18] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %common_high_snapshot_value_16, %common_high_snapshot_value_17, %common_high_snapshot_value_18 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } else { + %x_row_common_high_160 = index.mul %x_row_common_16, %x_row_stride : index + %x_row_common_high_161 = index.add %x_sequence_base, %x_row_common_high_160 : index + %x_index_common_high_160 = index.add %x_row_common_high_161, %channel : index + %x_index_common_high_16 = index.assume %x_index_common_high_160 [range(%x_index_common_high_160, 0, %flat_last)] : index + %common_high_raw_value_16 = vector.load %x_v[%x_index_common_high_16] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_high_170 = index.mul %x_row_common_17, %x_row_stride : index + %x_row_common_high_171 = index.add %x_sequence_base, %x_row_common_high_170 : index + %x_index_common_high_170 = index.add %x_row_common_high_171, %channel : index + %x_index_common_high_17 = index.assume %x_index_common_high_170 [range(%x_index_common_high_170, 0, %flat_last)] : index + %common_high_raw_value_17 = vector.load %x_v[%x_index_common_high_17] : view<1073741824xf32> -> vector<[%packet_width]xf32> + %x_row_common_high_180 = index.mul %x_row_common_18, %x_row_stride : index + %x_row_common_high_181 = index.add %x_sequence_base, %x_row_common_high_180 : index + %x_index_common_high_180 = index.add %x_row_common_high_181, %channel : index + %x_index_common_high_18 = index.assume %x_index_common_high_180 [range(%x_index_common_high_180, 0, %flat_last)] : index + %common_high_raw_value_18 = vector.load %x_v[%x_index_common_high_18] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %common_high_raw_value_16, %common_high_raw_value_17, %common_high_raw_value_18 : vector<[%packet_width]xf32>, vector<[%packet_width]xf32>, vector<[%packet_width]xf32> + } + // Filter and state are disjoint from x/dst. Each wave loads its lane's four + // filter values once and reuses them for 16 outputs. + %filter_byte0 = index.mul %channel, %c16 : index + %filter_byte = index.assume %filter_byte0 [range(%filter_byte0, 0, %filter_last), mul(%filter_byte0, %packet_bytes)] : index + %filter_offset = index.cast %filter_byte : index to offset + %filter_v = buffer.view %filt_na[%filter_offset] : buffer -> view<[%packet_width]x4xf32> + %filter_values = vector.load %filter_v[%c0, %c0] : view<[%packet_width]x4xf32> -> vector<[%packet_width]x4xf32> + %filter_transposed = vector.transpose<[1, 0]> %filter_values : vector<[%packet_width]x4xf32> -> vector<4x[%packet_width]xf32> + %w0 = vector.extract %filter_transposed[0] : vector<4x[%packet_width]xf32> -> vector<[%packet_width]xf32> + %w1 = vector.extract %filter_transposed[1] : vector<4x[%packet_width]xf32> -> vector<[%packet_width]xf32> + %w2 = vector.extract %filter_transposed[2] : vector<4x[%packet_width]xf32> -> vector<[%packet_width]xf32> + %w3 = vector.extract %filter_transposed[3] : vector<4x[%packet_width]xf32> -> vector<[%packet_width]xf32> + + %state0 = scf.if %first_wave -> (vector<[%packet_width]xf32>) { + %state_row_00 = index.mul %c0, %state_row_stride : index + %state_row_01 = index.add %state_sequence_base, %state_row_00 : index + %state_index_00 = index.add %state_row_01, %channel : index + %state_index_0 = index.assume %state_index_00 [range(%state_index_00, 0, %flat_last)] : index + %loaded = vector.load %state_v[%state_index_0] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %loaded : vector<[%packet_width]xf32> + } else { + scf.yield %c0_f32 : vector<[%packet_width]xf32> + } + + %state1 = scf.if %first_wave -> (vector<[%packet_width]xf32>) { + %state_row_10 = index.mul %c1, %state_row_stride : index + %state_row_11 = index.add %state_sequence_base, %state_row_10 : index + %state_index_10 = index.add %state_row_11, %channel : index + %state_index_1 = index.assume %state_index_10 [range(%state_index_10, 0, %flat_last)] : index + %loaded = vector.load %state_v[%state_index_1] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %loaded : vector<[%packet_width]xf32> + } else { + scf.yield %c0_f32 : vector<[%packet_width]xf32> + } + + %state2 = scf.if %first_wave -> (vector<[%packet_width]xf32>) { + %state_row_20 = index.mul %c2, %state_row_stride : index + %state_row_21 = index.add %state_sequence_base, %state_row_20 : index + %state_index_20 = index.add %state_row_21, %channel : index + %state_index_2 = index.assume %state_index_20 [range(%state_index_20, 0, %flat_last)] : index + %loaded = vector.load %state_v[%state_index_2] : view<1073741824xf32> -> vector<[%packet_width]xf32> + scf.yield %loaded : vector<[%packet_width]xf32> + } else { + scf.yield %c0_f32 : vector<[%packet_width]xf32> + } + + // Consume every X preload before the barrier because gfx11 s_barrier does not drain VMEM. + // Only the 16 activated values remain live across the rendezvous. + %v0_0 = scf.select %first_wave, %state0, %xr0 : vector<[%packet_width]xf32> + %v0_1 = scf.select %first_wave, %state1, %xr1 : vector<[%packet_width]xf32> + %v0_2 = scf.select %first_wave, %state2, %xr2 : vector<[%packet_width]xf32> + %p0_0 = vector.mulf %v0_0, %w0 : vector<[%packet_width]xf32> + %p0_1 = vector.fmaf %v0_1, %w1, %p0_0 : vector<[%packet_width]xf32> + %p0_2 = vector.fmaf %v0_2, %w2, %p0_1 : vector<[%packet_width]xf32> + %acc0 = vector.fmaf %xr3, %w3, %p0_2 : vector<[%packet_width]xf32> + %activated0 = vector.siluf %acc0 : vector<[%packet_width]xf32> + + %v1_0 = scf.select %first_wave, %state1, %xr1 : vector<[%packet_width]xf32> + %v1_1 = scf.select %first_wave, %state2, %xr2 : vector<[%packet_width]xf32> + %p1_0 = vector.mulf %v1_0, %w0 : vector<[%packet_width]xf32> + %p1_1 = vector.fmaf %v1_1, %w1, %p1_0 : vector<[%packet_width]xf32> + %p1_2 = vector.fmaf %xr3, %w2, %p1_1 : vector<[%packet_width]xf32> + %acc1 = vector.fmaf %xr4, %w3, %p1_2 : vector<[%packet_width]xf32> + %activated1 = vector.siluf %acc1 : vector<[%packet_width]xf32> + + %v2_0 = scf.select %first_wave, %state2, %xr2 : vector<[%packet_width]xf32> + %p2_0 = vector.mulf %v2_0, %w0 : vector<[%packet_width]xf32> + %p2_1 = vector.fmaf %xr3, %w1, %p2_0 : vector<[%packet_width]xf32> + %p2_2 = vector.fmaf %xr4, %w2, %p2_1 : vector<[%packet_width]xf32> + %acc2 = vector.fmaf %xr5, %w3, %p2_2 : vector<[%packet_width]xf32> + %activated2 = vector.siluf %acc2 : vector<[%packet_width]xf32> + + %p3_0 = vector.mulf %xr3, %w0 : vector<[%packet_width]xf32> + %p3_1 = vector.fmaf %xr4, %w1, %p3_0 : vector<[%packet_width]xf32> + %p3_2 = vector.fmaf %xr5, %w2, %p3_1 : vector<[%packet_width]xf32> + %acc3 = vector.fmaf %xr6, %w3, %p3_2 : vector<[%packet_width]xf32> + %activated3 = vector.siluf %acc3 : vector<[%packet_width]xf32> + + %p4_0 = vector.mulf %xr4, %w0 : vector<[%packet_width]xf32> + %p4_1 = vector.fmaf %xr5, %w1, %p4_0 : vector<[%packet_width]xf32> + %p4_2 = vector.fmaf %xr6, %w2, %p4_1 : vector<[%packet_width]xf32> + %acc4 = vector.fmaf %xr7, %w3, %p4_2 : vector<[%packet_width]xf32> + %activated4 = vector.siluf %acc4 : vector<[%packet_width]xf32> + + %p5_0 = vector.mulf %xr5, %w0 : vector<[%packet_width]xf32> + %p5_1 = vector.fmaf %xr6, %w1, %p5_0 : vector<[%packet_width]xf32> + %p5_2 = vector.fmaf %xr7, %w2, %p5_1 : vector<[%packet_width]xf32> + %acc5 = vector.fmaf %xr8, %w3, %p5_2 : vector<[%packet_width]xf32> + %activated5 = vector.siluf %acc5 : vector<[%packet_width]xf32> + + %p6_0 = vector.mulf %xr6, %w0 : vector<[%packet_width]xf32> + %p6_1 = vector.fmaf %xr7, %w1, %p6_0 : vector<[%packet_width]xf32> + %p6_2 = vector.fmaf %xr8, %w2, %p6_1 : vector<[%packet_width]xf32> + %acc6 = vector.fmaf %xr9, %w3, %p6_2 : vector<[%packet_width]xf32> + %activated6 = vector.siluf %acc6 : vector<[%packet_width]xf32> + + %p7_0 = vector.mulf %xr7, %w0 : vector<[%packet_width]xf32> + %p7_1 = vector.fmaf %xr8, %w1, %p7_0 : vector<[%packet_width]xf32> + %p7_2 = vector.fmaf %xr9, %w2, %p7_1 : vector<[%packet_width]xf32> + %acc7 = vector.fmaf %xr10, %w3, %p7_2 : vector<[%packet_width]xf32> + %activated7 = vector.siluf %acc7 : vector<[%packet_width]xf32> + + %p8_0 = vector.mulf %xr8, %w0 : vector<[%packet_width]xf32> + %p8_1 = vector.fmaf %xr9, %w1, %p8_0 : vector<[%packet_width]xf32> + %p8_2 = vector.fmaf %xr10, %w2, %p8_1 : vector<[%packet_width]xf32> + %acc8 = vector.fmaf %xr11, %w3, %p8_2 : vector<[%packet_width]xf32> + %activated8 = vector.siluf %acc8 : vector<[%packet_width]xf32> + + %p9_0 = vector.mulf %xr9, %w0 : vector<[%packet_width]xf32> + %p9_1 = vector.fmaf %xr10, %w1, %p9_0 : vector<[%packet_width]xf32> + %p9_2 = vector.fmaf %xr11, %w2, %p9_1 : vector<[%packet_width]xf32> + %acc9 = vector.fmaf %xr12, %w3, %p9_2 : vector<[%packet_width]xf32> + %activated9 = vector.siluf %acc9 : vector<[%packet_width]xf32> + + %p10_0 = vector.mulf %xr10, %w0 : vector<[%packet_width]xf32> + %p10_1 = vector.fmaf %xr11, %w1, %p10_0 : vector<[%packet_width]xf32> + %p10_2 = vector.fmaf %xr12, %w2, %p10_1 : vector<[%packet_width]xf32> + %acc10 = vector.fmaf %xr13, %w3, %p10_2 : vector<[%packet_width]xf32> + %activated10 = vector.siluf %acc10 : vector<[%packet_width]xf32> + + %p11_0 = vector.mulf %xr11, %w0 : vector<[%packet_width]xf32> + %p11_1 = vector.fmaf %xr12, %w1, %p11_0 : vector<[%packet_width]xf32> + %p11_2 = vector.fmaf %xr13, %w2, %p11_1 : vector<[%packet_width]xf32> + %acc11 = vector.fmaf %xr14, %w3, %p11_2 : vector<[%packet_width]xf32> + %activated11 = vector.siluf %acc11 : vector<[%packet_width]xf32> + + %p12_0 = vector.mulf %xr12, %w0 : vector<[%packet_width]xf32> + %p12_1 = vector.fmaf %xr13, %w1, %p12_0 : vector<[%packet_width]xf32> + %p12_2 = vector.fmaf %xr14, %w2, %p12_1 : vector<[%packet_width]xf32> + %acc12 = vector.fmaf %xr15, %w3, %p12_2 : vector<[%packet_width]xf32> + %activated12 = vector.siluf %acc12 : vector<[%packet_width]xf32> + + %p13_0 = vector.mulf %xr13, %w0 : vector<[%packet_width]xf32> + %p13_1 = vector.fmaf %xr14, %w1, %p13_0 : vector<[%packet_width]xf32> + %p13_2 = vector.fmaf %xr15, %w2, %p13_1 : vector<[%packet_width]xf32> + %acc13 = vector.fmaf %xr16, %w3, %p13_2 : vector<[%packet_width]xf32> + %activated13 = vector.siluf %acc13 : vector<[%packet_width]xf32> + + %p14_0 = vector.mulf %xr14, %w0 : vector<[%packet_width]xf32> + %p14_1 = vector.fmaf %xr15, %w1, %p14_0 : vector<[%packet_width]xf32> + %p14_2 = vector.fmaf %xr16, %w2, %p14_1 : vector<[%packet_width]xf32> + %acc14 = vector.fmaf %xr17, %w3, %p14_2 : vector<[%packet_width]xf32> + %activated14 = vector.siluf %acc14 : vector<[%packet_width]xf32> + + %p15_0 = vector.mulf %xr15, %w0 : vector<[%packet_width]xf32> + %p15_1 = vector.fmaf %xr16, %w1, %p15_0 : vector<[%packet_width]xf32> + %p15_2 = vector.fmaf %xr17, %w2, %p15_1 : vector<[%packet_width]xf32> + %acc15 = vector.fmaf %xr18, %w3, %p15_2 : vector<[%packet_width]xf32> + %activated15 = vector.siluf %acc15 : vector<[%packet_width]xf32> + + // Uniform load/compute-before-store boundary. There is no LDS allocation + // and no destination store above this point. + kernel.barrier scope(workgroup) ordering(acq_rel) + + %dst_sequence_span0 = index.mul %n_t, %d_inner : index + %dst_sequence_span = index.assume %dst_sequence_span0 [range(%dst_sequence_span0, 4194304, 5242880), mul(%dst_sequence_span0, 16384)] : index + %dst_sequence_base0 = index.mul %sequence, %dst_sequence_span : index + %dst_sequence_base = index.assume %dst_sequence_base0 [range(%dst_sequence_base0, 0, %flat_last)] : index + + %token_0_raw = index.add %wave_token_base, %c0 : index + %token_0 = index.assume %token_0_raw [range(%token_0_raw, 0, 496)] : index + %dst_row_0_raw = index.mul %token_0, %d_inner : index + %dst_row_0 = index.add %dst_sequence_base, %dst_row_0_raw : index + %dst_index_0_raw = index.add %dst_row_0, %channel : index + %dst_index_0 = index.assume %dst_index_0_raw [range(%dst_index_0_raw, 0, %flat_last)] : index + vector.store %activated0, %dst_v[%dst_index_0] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_1_raw = index.add %wave_token_base, %c1 : index + %token_1 = index.assume %token_1_raw [range(%token_1_raw, 1, 497)] : index + %dst_row_1_raw = index.mul %token_1, %d_inner : index + %dst_row_1 = index.add %dst_sequence_base, %dst_row_1_raw : index + %dst_index_1_raw = index.add %dst_row_1, %channel : index + %dst_index_1 = index.assume %dst_index_1_raw [range(%dst_index_1_raw, 0, %flat_last)] : index + vector.store %activated1, %dst_v[%dst_index_1] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_2_raw = index.add %wave_token_base, %c2 : index + %token_2 = index.assume %token_2_raw [range(%token_2_raw, 2, 498)] : index + %dst_row_2_raw = index.mul %token_2, %d_inner : index + %dst_row_2 = index.add %dst_sequence_base, %dst_row_2_raw : index + %dst_index_2_raw = index.add %dst_row_2, %channel : index + %dst_index_2 = index.assume %dst_index_2_raw [range(%dst_index_2_raw, 0, %flat_last)] : index + vector.store %activated2, %dst_v[%dst_index_2] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_3_raw = index.add %wave_token_base, %c3 : index + %token_3 = index.assume %token_3_raw [range(%token_3_raw, 3, 499)] : index + %dst_row_3_raw = index.mul %token_3, %d_inner : index + %dst_row_3 = index.add %dst_sequence_base, %dst_row_3_raw : index + %dst_index_3_raw = index.add %dst_row_3, %channel : index + %dst_index_3 = index.assume %dst_index_3_raw [range(%dst_index_3_raw, 0, %flat_last)] : index + vector.store %activated3, %dst_v[%dst_index_3] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_4_raw = index.add %wave_token_base, %c4 : index + %token_4 = index.assume %token_4_raw [range(%token_4_raw, 4, 500)] : index + %dst_row_4_raw = index.mul %token_4, %d_inner : index + %dst_row_4 = index.add %dst_sequence_base, %dst_row_4_raw : index + %dst_index_4_raw = index.add %dst_row_4, %channel : index + %dst_index_4 = index.assume %dst_index_4_raw [range(%dst_index_4_raw, 0, %flat_last)] : index + vector.store %activated4, %dst_v[%dst_index_4] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_5_raw = index.add %wave_token_base, %c5 : index + %token_5 = index.assume %token_5_raw [range(%token_5_raw, 5, 501)] : index + %dst_row_5_raw = index.mul %token_5, %d_inner : index + %dst_row_5 = index.add %dst_sequence_base, %dst_row_5_raw : index + %dst_index_5_raw = index.add %dst_row_5, %channel : index + %dst_index_5 = index.assume %dst_index_5_raw [range(%dst_index_5_raw, 0, %flat_last)] : index + vector.store %activated5, %dst_v[%dst_index_5] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_6_raw = index.add %wave_token_base, %c6 : index + %token_6 = index.assume %token_6_raw [range(%token_6_raw, 6, 502)] : index + %dst_row_6_raw = index.mul %token_6, %d_inner : index + %dst_row_6 = index.add %dst_sequence_base, %dst_row_6_raw : index + %dst_index_6_raw = index.add %dst_row_6, %channel : index + %dst_index_6 = index.assume %dst_index_6_raw [range(%dst_index_6_raw, 0, %flat_last)] : index + vector.store %activated6, %dst_v[%dst_index_6] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_7_raw = index.add %wave_token_base, %c7 : index + %token_7 = index.assume %token_7_raw [range(%token_7_raw, 7, 503)] : index + %dst_row_7_raw = index.mul %token_7, %d_inner : index + %dst_row_7 = index.add %dst_sequence_base, %dst_row_7_raw : index + %dst_index_7_raw = index.add %dst_row_7, %channel : index + %dst_index_7 = index.assume %dst_index_7_raw [range(%dst_index_7_raw, 0, %flat_last)] : index + vector.store %activated7, %dst_v[%dst_index_7] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_8_raw = index.add %wave_token_base, %c8 : index + %token_8 = index.assume %token_8_raw [range(%token_8_raw, 8, 504)] : index + %dst_row_8_raw = index.mul %token_8, %d_inner : index + %dst_row_8 = index.add %dst_sequence_base, %dst_row_8_raw : index + %dst_index_8_raw = index.add %dst_row_8, %channel : index + %dst_index_8 = index.assume %dst_index_8_raw [range(%dst_index_8_raw, 0, %flat_last)] : index + vector.store %activated8, %dst_v[%dst_index_8] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_9_raw = index.add %wave_token_base, %c9 : index + %token_9 = index.assume %token_9_raw [range(%token_9_raw, 9, 505)] : index + %dst_row_9_raw = index.mul %token_9, %d_inner : index + %dst_row_9 = index.add %dst_sequence_base, %dst_row_9_raw : index + %dst_index_9_raw = index.add %dst_row_9, %channel : index + %dst_index_9 = index.assume %dst_index_9_raw [range(%dst_index_9_raw, 0, %flat_last)] : index + vector.store %activated9, %dst_v[%dst_index_9] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_10_raw = index.add %wave_token_base, %c10 : index + %token_10 = index.assume %token_10_raw [range(%token_10_raw, 10, 506)] : index + %dst_row_10_raw = index.mul %token_10, %d_inner : index + %dst_row_10 = index.add %dst_sequence_base, %dst_row_10_raw : index + %dst_index_10_raw = index.add %dst_row_10, %channel : index + %dst_index_10 = index.assume %dst_index_10_raw [range(%dst_index_10_raw, 0, %flat_last)] : index + vector.store %activated10, %dst_v[%dst_index_10] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_11_raw = index.add %wave_token_base, %c11 : index + %token_11 = index.assume %token_11_raw [range(%token_11_raw, 11, 507)] : index + %dst_row_11_raw = index.mul %token_11, %d_inner : index + %dst_row_11 = index.add %dst_sequence_base, %dst_row_11_raw : index + %dst_index_11_raw = index.add %dst_row_11, %channel : index + %dst_index_11 = index.assume %dst_index_11_raw [range(%dst_index_11_raw, 0, %flat_last)] : index + vector.store %activated11, %dst_v[%dst_index_11] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_12_raw = index.add %wave_token_base, %c12 : index + %token_12 = index.assume %token_12_raw [range(%token_12_raw, 12, 508)] : index + %dst_row_12_raw = index.mul %token_12, %d_inner : index + %dst_row_12 = index.add %dst_sequence_base, %dst_row_12_raw : index + %dst_index_12_raw = index.add %dst_row_12, %channel : index + %dst_index_12 = index.assume %dst_index_12_raw [range(%dst_index_12_raw, 0, %flat_last)] : index + vector.store %activated12, %dst_v[%dst_index_12] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_13_raw = index.add %wave_token_base, %c13 : index + %token_13 = index.assume %token_13_raw [range(%token_13_raw, 13, 509)] : index + %dst_row_13_raw = index.mul %token_13, %d_inner : index + %dst_row_13 = index.add %dst_sequence_base, %dst_row_13_raw : index + %dst_index_13_raw = index.add %dst_row_13, %channel : index + %dst_index_13 = index.assume %dst_index_13_raw [range(%dst_index_13_raw, 0, %flat_last)] : index + vector.store %activated13, %dst_v[%dst_index_13] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_14_raw = index.add %wave_token_base, %c14 : index + %token_14 = index.assume %token_14_raw [range(%token_14_raw, 14, 510)] : index + %dst_row_14_raw = index.mul %token_14, %d_inner : index + %dst_row_14 = index.add %dst_sequence_base, %dst_row_14_raw : index + %dst_index_14_raw = index.add %dst_row_14, %channel : index + %dst_index_14 = index.assume %dst_index_14_raw [range(%dst_index_14_raw, 0, %flat_last)] : index + vector.store %activated14, %dst_v[%dst_index_14] : vector<[%packet_width]xf32>, view<1073741824xf32> + %token_15_raw = index.add %wave_token_base, %c15 : index + %token_15 = index.assume %token_15_raw [range(%token_15_raw, 15, 511)] : index + %dst_row_15_raw = index.mul %token_15, %d_inner : index + %dst_row_15 = index.add %dst_sequence_base, %dst_row_15_raw : index + %dst_index_15_raw = index.add %dst_row_15, %channel : index + %dst_index_15 = index.assume %dst_index_15_raw [range(%dst_index_15_raw, 0, %flat_last)] : index + vector.store %activated15, %dst_v[%dst_index_15] : vector<[%packet_width]xf32>, view<1073741824xf32> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/rollback_store.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/rollback_store.loom new file mode 100644 index 000000000000..31aa9aea96f6 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/rollback_store.loom @@ -0,0 +1,77 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.rollback_store_carried(%slot: index, %cache_count: index, %n_t: index, %cache_sequence_base: index, %channel: index, %old0: f32, %old1: f32, %old2: f32, %cache: buffer) + +template.decl @llm.ssm_conv.rollback_store_x(%slot: index, %cache_count: index, %n_t: index, %cache_sequence_base: index, %token: index, %channel: index, %value: f32, %cache: buffer) + +template.def<@llm.ssm_conv.rollback_store_carried> device @llm_ssm_conv_rollback_store_carried(%slot: index, %cache_count: index, %n_t: index, %cache_sequence_base: index, %channel: index, %old0: f32, %old1: f32, %old2: f32, %cache: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %active = index.cmp ult, %slot, %cache_count : index + scf.if %active { + %has_full_history = index.cmp ult, %slot, %n_t : index + %start = scf.if %has_full_history -> (index) { + %candidate = index.sub %n_t, %slot : index + scf.yield %candidate : index + } else { + scf.yield %c0 : index + } + %cache_g = buffer.assume.memory_space %cache : buffer + %cache_v = buffer.view %cache_g[%base] : buffer -> view<589824xf32> + scf.for %row = [%c0 to %c3 step %c1] { + %source_row = index.add %start, %row : index + %is_carried = index.cmp ult, %source_row, %c3 : index + scf.if %is_carried { + %is_zero = index.cmp eq, %source_row, %c0 : index + %is_one = index.cmp eq, %source_row, %c1 : index + %old12 = scf.select %is_one, %old1, %old2 : f32 + %value = scf.select %is_zero, %old0, %old12 : f32 + %channel_base0 = index.mul %channel, %c3 : index + %channel_base = index.assume %channel_base0 [range(%channel_base0, 0, 196605), mul(%channel_base0, 3)] : index + %cache_channel_index0 = index.add %channel_base, %row : index + %cache_channel_index = index.assume %cache_channel_index0 [range(%cache_channel_index0, 0, 196607)] : index + %cache_index0 = index.add %cache_sequence_base, %cache_channel_index : index + %cache_index = index.assume %cache_index0 [range(%cache_index0, 0, 589823)] : index + view.store %value, %cache_v[%cache_index] : f32, view<589824xf32> + } + } + } + template.return +} +template.def<@llm.ssm_conv.rollback_store_x> device @llm_ssm_conv_rollback_store_x(%slot: index, %cache_count: index, %n_t: index, %cache_sequence_base: index, %token: index, %channel: index, %value: f32, %cache: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c3 = index.constant 3 : index + %active = index.cmp ult, %slot, %cache_count : index + %has_full_history = index.cmp ult, %slot, %n_t : index + %start = scf.if %has_full_history -> (index) { + %candidate = index.sub %n_t, %slot : index + scf.yield %candidate : index + } else { + scf.yield %c0 : index + } + %source_row = index.add %token, %c3 : index + %end = index.add %start, %c3 : index + %at_or_after_start = index.cmp uge, %source_row, %start : index + %before_end = index.cmp ult, %source_row, %end : index + %within_window = scalar.andi %at_or_after_start, %before_end : i1 + %writes = scalar.andi %active, %within_window : i1 + scf.if %writes { + %row0 = index.sub %source_row, %start : index + %row = index.assume %row0 [range(%row0, 0, 2)] : index + %channel_base0 = index.mul %channel, %c3 : index + %channel_base = index.assume %channel_base0 [range(%channel_base0, 0, 196605), mul(%channel_base0, 3)] : index + %cache_channel_index0 = index.add %channel_base, %row : index + %cache_channel_index = index.assume %cache_channel_index0 [range(%cache_channel_index0, 0, 196607)] : index + %cache_index0 = index.add %cache_sequence_base, %cache_channel_index : index + %cache_index = index.assume %cache_index0 [range(%cache_index0, 0, 589823)] : index + %cache_g = buffer.assume.memory_space %cache : buffer + %cache_v = buffer.view %cache_g[%base] : buffer -> view<589824xf32> + view.store %value, %cache_v[%cache_index] : f32, view<589824xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/snapshot_window_tail.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/snapshot_window_tail.loom new file mode 100644 index 000000000000..e04ef39c6714 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/snapshot_window_tail.loom @@ -0,0 +1,147 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.snapshot_window_tail(%state: buffer, %x: buffer, %dst: buffer, %cache: buffer) + +config.decl @llm.ssm_conv.snapshot.d_conv : %value: index where [range(%value, 4, 4)] + +config.decl @llm.ssm_conv.snapshot.d_inner : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +config.decl @llm.ssm_conv.snapshot.n_t : %value: index where [range(%value, 512, 512)] + +config.decl @llm.ssm_conv.snapshot.n_s : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.snapshot.state_row_stride : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.snapshot.state_channel_stride : %value: index where [range(%value, 3, 3)] + +config.decl @llm.ssm_conv.snapshot.x_row_stride : %value: index where [range(%value, 1, 16777216)] + +config.decl @llm.ssm_conv.snapshot.dst_row_stride : %value: index where [range(%value, 1, 16777216)] + +config.decl @llm.ssm_conv.snapshot.cache_row_stride : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.snapshot.cache_channel_stride : %value: index where [range(%value, 3, 3)] + +config.decl @llm.ssm_conv.snapshot.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] + +template.def<@llm.ssm_conv.snapshot_window_tail> device @llm_ssm_conv_snapshot_window_tail(%state: buffer, %x: buffer, %dst: buffer, %cache: buffer) { + %base = index.constant 0 : offset + %cneg1 = index.constant -1 : index + %x_snapshot_rows = index.constant 61 : index + + %d_conv = config.get @llm.ssm_conv.snapshot.d_conv : index + %d_inner = config.get @llm.ssm_conv.snapshot.d_inner : index + %n_t = config.get @llm.ssm_conv.snapshot.n_t : index + %n_s = config.get @llm.ssm_conv.snapshot.n_s : index + %srs = config.get @llm.ssm_conv.snapshot.state_row_stride : index + %scs = config.get @llm.ssm_conv.snapshot.state_channel_stride : index + %xrs = config.get @llm.ssm_conv.snapshot.x_row_stride : index + %drs = config.get @llm.ssm_conv.snapshot.dst_row_stride : index + %crs = config.get @llm.ssm_conv.snapshot.cache_row_stride : index + %ccs = config.get @llm.ssm_conv.snapshot.cache_channel_stride : index + %wg = config.get @llm.ssm_conv.snapshot.workgroup_size : index + + %di = index.assume %d_inner [range(%d_inner, 8192, 10240), mul(%d_inner, 32)] : index + %nt = index.assume %n_t [range(%n_t, 512, 512)] : index + %ns = index.assume %n_s [range(%n_s, 1, 1)] : index + %srsb = index.assume %srs [range(%srs, 1, 1)] : index + %scsb = index.assume %scs [range(%scs, 3, 3)] : index + %xrsb = index.assume %xrs [range(%xrs, 1, 16777216)] : index + %drsb = index.assume %drs [range(%drs, 1, 16777216)] : index + %crsb = index.assume %crs [range(%crs, 1, 1)] : index + %ccsb = index.assume %ccs [range(%ccs, 3, 3)] : index + %carry0 = index.add %d_conv, %cneg1 : index + %carry = index.assume %carry0 [range(%carry0, 3, 3)] : index + %snapshot_rows0 = index.add %carry, %x_snapshot_rows : index + %snapshot_rows = index.assume %snapshot_rows0 [range(%snapshot_rows0, 64, 64)] : index + + %rows0 = index.mul %snapshot_rows, %ns : index + %total0 = index.mul %rows0, %di : index + %total = index.assume %total0 [range(%total0, 1, 1073741823)] : index + + %group0 = kernel.workgroup.id : index + %lane0 = kernel.workitem.id : index + %group = index.assume %group0 [range(%group0, 0, 16777215)] : index + %lane = index.assume %lane0 [range(%lane0, 0, 1023)] : index + %off0 = index.mul %group, %wg : index + %idx0 = index.add %off0, %lane : index + %idx = index.assume %idx0 [range(%idx0, 0, 1073741823)] : index + %in_bounds = index.cmp ult, %idx, %total : index + + %state_g = buffer.assume.memory_space %state : buffer + %x_g = buffer.assume.memory_space %x : buffer + %dst_g = buffer.assume.memory_space %dst : buffer + %c_g = buffer.assume.memory_space %cache : buffer + %state_na, %x_na, %dst_na, %c_na = buffer.assume.noalias %state_g, %x_g, %dst_g, %c_g : buffer, buffer, buffer, buffer + %state_view = buffer.view %state_na[%base] : buffer -> view<1073741824xf32> + %x_view = buffer.view %x_na[%base] : buffer -> view<1073741824xf32> + %dst_view = buffer.view %dst_na[%base] : buffer -> view<1073741824xf32> + %c_view = buffer.view %c_na[%base] : buffer -> view<1073741824xf32> + + scf.if %in_bounds { + // Consecutive lanes take consecutive channels, so both sides are contiguous + // bursts across the wave. + %chan = index.rem %idx, %di : index + %flat_row = index.div %idx, %di : index + %snapshot_row = index.rem %flat_row, %snapshot_rows : index + %seq = index.div %flat_row, %snapshot_rows : index + + // All 64 snapshot rows are written to their natural CONCAT row. + %d_win0 = index.add %nt, %carry : index + %d_span = index.mul %d_win0, %drsb : index + %d_seq = index.mul %seq, %d_span : index + %oi0 = index.mul %snapshot_row, %drsb : index + %oi1 = index.add %oi0, %d_seq : index + %oi2 = index.add %oi1, %chan : index + %oi = index.assume %oi2 [range(%oi2, 0, 1073741823)] : index + + %is_state = index.cmp ult, %snapshot_row, %carry : index + %snapshot = scf.if %is_state -> (f32) { + // Rows 0..2 snapshot the old recurrent state. In the same workitem, + // update the corresponding persistent-cache row from x rows 509..511. + %state_span = index.mul %di, %scsb : index + %state_seq = index.mul %seq, %state_span : index + %si0 = index.mul %snapshot_row, %srsb : index + %si1 = index.add %si0, %state_seq : index + %si_chan = index.mul %chan, %scsb : index + %si2 = index.add %si1, %si_chan : index + %si = index.assume %si2 [range(%si2, 0, 1073741823)] : index + + %x_tail0 = index.sub %nt, %carry : index + %x_tail = index.add %x_tail0, %snapshot_row : index + %x_span = index.mul %nt, %xrsb : index + %x_seq = index.mul %seq, %x_span : index + %xi0 = index.mul %x_tail, %xrsb : index + %xi1 = index.add %xi0, %x_seq : index + %xi2 = index.add %xi1, %chan : index + %xi = index.assume %xi2 [range(%xi2, 0, 1073741823)] : index + + %cache_span = index.mul %di, %ccsb : index + %cache_seq = index.mul %seq, %cache_span : index + %ci0 = index.mul %snapshot_row, %crsb : index + %ci1 = index.add %ci0, %cache_seq : index + %ci_chan = index.mul %chan, %ccsb : index + %ci2 = index.add %ci1, %ci_chan : index + %ci = index.assume %ci2 [range(%ci2, 0, 1073741823)] : index + + %old_state = view.load %state_view[%si] : view<1073741824xf32> -> f32 + %new_state = view.load %x_view[%xi] : view<1073741824xf32> -> f32 + view.store %new_state, %c_view[%ci] : f32, view<1073741824xf32> + scf.yield %old_state : f32 + } else { + // Rows 3..63 snapshot x rows 0..60 before their arena span is reused. + %x_row = index.sub %snapshot_row, %carry : index + %x_span = index.mul %nt, %xrsb : index + %x_seq = index.mul %seq, %x_span : index + %xi0 = index.mul %x_row, %xrsb : index + %xi1 = index.add %xi0, %x_seq : index + %xi2 = index.add %xi1, %chan : index + %xi = index.assume %xi2 [range(%xi2, 0, 1073741823)] : index + %x_value = view.load %x_view[%xi] : view<1073741824xf32> -> f32 + scf.yield %x_value : f32 + } + view.store %snapshot, %dst_view[%oi] : f32, view<1073741824xf32> + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/state_materialize.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/state_materialize.loom new file mode 100644 index 000000000000..f7e9d3cc2d0f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/ssm_conv/state_materialize.loom @@ -0,0 +1,176 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.rollback_store_carried(%slot: index, %cache_count: index, %n_t: index, %cache_sequence_base: index, %channel: index, %old0: f32, %old1: f32, %old2: f32, %cache: buffer) + +template.decl @llm.ssm_conv.rollback_store_x(%slot: index, %cache_count: index, %n_t: index, %cache_sequence_base: index, %token: index, %channel: index, %value: f32, %cache: buffer) + +template.decl @llm.ssm_conv.decode_update(%prev0: f32, %prev1: f32, %prev2: f32, %new_x: f32, %w0: f32, %w1: f32, %w2: f32, %w3: f32) -> (f32) + +template.decl @llm.ssm_conv.state_materialized_body(%d_inner0: index, %n_t0: index, %n_s0: index, %cache_count0: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) + +template.def<@llm.ssm_conv.state_materialized_body> device @llm_ssm_conv_state_materialized_body(%d_inner0: index, %n_t0: index, %n_s0: index, %cache_count0: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) { + %base = index.constant 0 : offset + %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 + %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, 3)] : index + %cache_count = index.assume %cache_count0 [range(%cache_count0, 1, 5)] : index + %group0 = kernel.workgroup.id : index + %sequence0 = kernel.workgroup.id : index + %lane0 = kernel.workitem.id : index + %group = index.assume %group0 [range(%group0, 0, 255)] : index + %sequence, %launch_n_s = index.assume %sequence0, %n_s [lt(%sequence0, %n_s)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index + %c256 = index.constant 256 : index + %channel_base = index.mul %group, %c256 : index + %channel0 = index.add %channel_base, %lane : index + %channel = index.assume %channel0 [range(%channel0, 0, 65535)] : index + %in_bounds = index.cmp ult, %channel, %d_inner : index + + %state_g = buffer.assume.memory_space %state : buffer + %x_g = buffer.assume.memory_space %x : buffer + %filt_g = buffer.assume.memory_space %filt : buffer + %dst_g = buffer.assume.memory_space %dst : buffer + %state_na, %x_na, %filt_na = buffer.assume.noalias %state_g, %x_g, %filt_g : buffer, buffer, buffer + %state_v = buffer.view %state_na[%base] : buffer -> view<589824xf32> + %x_v = buffer.view %x_na[%base] : buffer -> view<100663296xf32> + %filt_v = buffer.view %filt_na[%base] : buffer -> view<262144xf32> + %dst_v = buffer.view %dst_g[%base] : buffer -> view<100663296xf32> + + %state_sequence_span0 = index.mul %d_inner, %c3 : index + %state_sequence_span = index.assume %state_sequence_span0 [range(%state_sequence_span0, 96, 196608), mul(%state_sequence_span0, 96)] : index + %state_sequence_base0 = index.mul %sequence, %state_sequence_span : index + %state_sequence_base = index.assume %state_sequence_base0 [range(%state_sequence_base0, 0, 393216)] : index + %x_sequence_span0 = index.mul %n_t, %d_inner : index + %x_sequence_span = index.assume %x_sequence_span0 [range(%x_sequence_span0, 32, 33554432), mul(%x_sequence_span0, 32)] : index + %x_sequence_base0 = index.mul %sequence, %x_sequence_span : index + %x_sequence_base = index.assume %x_sequence_base0 [range(%x_sequence_base0, 0, 67108864)] : index + + scf.if %in_bounds { + %state_channel0 = index.mul %channel, %c3 : index + %state_channel = index.assume %state_channel0 [range(%state_channel0, 0, 196605), mul(%state_channel0, 3)] : index + %state_row0 = index.add %state_sequence_base, %state_channel : index + %state_row = index.assume %state_row0 [range(%state_row0, 0, 589821)] : index + %state_row1_0 = index.add %state_row, %c1 : index + %state_row1 = index.assume %state_row1_0 [range(%state_row1_0, 1, 589822)] : index + %state_row2_0 = index.add %state_row, %c2 : index + %state_row2 = index.assume %state_row2_0 [range(%state_row2_0, 2, 589823)] : index + + %old0 = view.load %state_v[%state_row] : view<589824xf32> -> f32 + %old1 = view.load %state_v[%state_row1] : view<589824xf32> -> f32 + %old2 = view.load %state_v[%state_row2] : view<589824xf32> -> f32 + + %filter_row0 = index.mul %channel, %c4 : index + %filter_row = index.assume %filter_row0 [range(%filter_row0, 0, 262140), mul(%filter_row0, 4)] : index + %weights = vector.load %filt_v[%filter_row] : view<262144xf32> -> vector<4xf32> + %w0 = vector.extract %weights[%c0] : vector<4xf32> -> f32 + %w1 = vector.extract %weights[%c1] : vector<4xf32> -> f32 + %w2 = vector.extract %weights[%c2] : vector<4xf32> -> f32 + %w3 = vector.extract %weights[%c3] : vector<4xf32> -> f32 + + template.apply<@llm.ssm_conv.rollback_store_carried>(%c0, %cache_count, %n_t, %state_sequence_base, %channel, %old0, %old1, %old2, %cache0) : (index, index, index, index, index, f32, f32, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_carried>(%c1, %cache_count, %n_t, %state_sequence_base, %channel, %old0, %old1, %old2, %cache1) : (index, index, index, index, index, f32, f32, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_carried>(%c2, %cache_count, %n_t, %state_sequence_base, %channel, %old0, %old1, %old2, %cache2) : (index, index, index, index, index, f32, f32, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_carried>(%c3, %cache_count, %n_t, %state_sequence_base, %channel, %old0, %old1, %old2, %cache3) : (index, index, index, index, index, f32, f32, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_carried>(%c4, %cache_count, %n_t, %state_sequence_base, %channel, %old0, %old1, %old2, %cache4) : (index, index, index, index, index, f32, f32, f32, buffer) + + // Replace the full three-value history at each rolled loop boundary. + // Keep only three steps and the final partial block unrolled. + %full_trip_count = index.div %n_t, %c3 : index + %full_token_count = index.mul %full_trip_count, %c3 : index + %trip0, %trip1, %trip2 = scf.for %block = [%c0 to %full_token_count step %c3](%block0 = %old0 : f32, %block1 = %old1 : f32, %block2 = %old2 : f32) -> (f32, f32, f32) { + %step0, %step1, %step2 = scf.for %phase = [%c0 to %c3 step %c1](%prev0 = %block0 : f32, %prev1 = %block1 : f32, %prev2 = %block2 : f32) -> (f32, f32, f32) unroll { + %token = index.add %block, %phase : index + + %row_base0 = index.mul %token, %d_inner : index + %row_base = index.assume %row_base0 [range(%row_base0, 0, 33488896), mul(%row_base0, 32)] : index + %row_index0 = index.add %row_base, %channel : index + %row_index = index.assume %row_index0 [range(%row_index0, 0, 33554431)] : index + %index0 = index.add %x_sequence_base, %row_index : index + %index = index.assume %index0 [range(%index0, 0, 100663295)] : index + %new_x = view.load %x_v[%index] : view<100663296xf32> -> f32 + %activated = template.apply<@llm.ssm_conv.decode_update>(%prev0, %prev1, %prev2, %new_x, %w0, %w1, %w2, %w3) : (f32, f32, f32, f32, f32, f32, f32, f32) -> (f32) + + template.apply<@llm.ssm_conv.rollback_store_x>(%c0, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache0) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c1, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache1) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c2, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache2) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c3, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache3) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c4, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache4) : (index, index, index, index, index, index, f32, buffer) + view.store %activated, %dst_v[%index] : f32, view<100663296xf32> + scf.yield %prev1, %prev2, %new_x : f32, f32, f32 + } + scf.yield %step0, %step1, %step2 : f32, f32, f32 + } + %final0, %final1, %final2 = scf.for %token = [%full_token_count to %n_t step %c1](%prev0 = %trip0 : f32, %prev1 = %trip1 : f32, %prev2 = %trip2 : f32) -> (f32, f32, f32) unroll { + %row_base0 = index.mul %token, %d_inner : index + %row_base = index.assume %row_base0 [range(%row_base0, 0, 33488896), mul(%row_base0, 32)] : index + %row_index0 = index.add %row_base, %channel : index + %row_index = index.assume %row_index0 [range(%row_index0, 0, 33554431)] : index + %index0 = index.add %x_sequence_base, %row_index : index + %index = index.assume %index0 [range(%index0, 0, 100663295)] : index + %new_x = view.load %x_v[%index] : view<100663296xf32> -> f32 + %activated = template.apply<@llm.ssm_conv.decode_update>(%prev0, %prev1, %prev2, %new_x, %w0, %w1, %w2, %w3) : (f32, f32, f32, f32, f32, f32, f32, f32) -> (f32) + + template.apply<@llm.ssm_conv.rollback_store_x>(%c0, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache0) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c1, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache1) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c2, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache2) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c3, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache3) : (index, index, index, index, index, index, f32, buffer) + template.apply<@llm.ssm_conv.rollback_store_x>(%c4, %cache_count, %n_t, %state_sequence_base, %token, %channel, %new_x, %cache4) : (index, index, index, index, index, index, f32, buffer) + view.store %activated, %dst_v[%index] : f32, view<100663296xf32> + scf.yield %prev1, %prev2, %new_x : f32, f32, f32 + } + } + template.return +} + +// Regression: distinct initial history must survive every rolled boundary. +// With history [1, 2, 3], five input tokens equal to 4, and filter +// [1, 2, 3, 4], the convolution is [30, 36, 39, 40, 40]. +// SiLU rounds each of these positive values back to itself in F32. +kernel.def @ssm_conv_rolling_history_check() { + %one = index.constant 1 : index + %wg = index.constant 256 : index + kernel.launch.config workgroups(%one, %one, %one) workgroup_size(%wg, %one, %one) : index +} launch(%state: buffer, %input: buffer, %filter: buffer, %output: buffer, %cache: buffer) { + %channels = index.constant 32 : index + %tokens = index.constant 5 : index + %one = index.constant 1 : index + template.apply<@llm.ssm_conv.state_materialized_body>(%channels, %tokens, %one, %one, %state, %input, %filter, %output, %cache, %cache, %cache, %cache, %cache) : (index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +check.case public @ssm_conv_rolling_history { + %state = check.generate.iota offset(1.0) step(1.0) period(3) : tensor<96xf32> + %input = check.generate.fill value(4.0) : tensor<160xf32> + %filter = check.generate.iota offset(1.0) step(1.0) period(4) : tensor<128xf32> + %separate = check.generate.fill value(-999.0) : tensor<160xf32> + %inplace = check.generate.fill value(-999.0) : tensor<160xf32> + %cache = check.generate.fill value(-999.0) : tensor<96xf32> + %expected_cache = check.generate.fill value(4.0) : tensor<96xf32> + kernel.launch @ssm_conv_rolling_history_check[](%state, %input, %filter, %separate, %cache) : [](tensor<96xf32>, tensor<160xf32>, tensor<128xf32>, tensor<160xf32>, tensor<96xf32>) + kernel.launch @ssm_conv_rolling_history_check[](%state, %input, %filter, %inplace, %state) : [](tensor<96xf32>, tensor<160xf32>, tensor<128xf32>, tensor<160xf32>, tensor<96xf32>) + check.expect.equal actual(%cache) expected(%expected_cache) : tensor<96xf32> + check.expect.equal actual(%state) expected(%expected_cache) : tensor<96xf32> + check.expect.equal actual(%inplace) expected(%separate) : tensor<160xf32> + %row0 = check.tensor.view %separate offset(0) : tensor<160xf32> -> tensor<32xf32> + %expected0 = check.generate.fill value(30.0) : tensor<32xf32> + check.expect.equal actual(%row0) expected(%expected0) : tensor<32xf32> + %row1 = check.tensor.view %separate offset(128) : tensor<160xf32> -> tensor<32xf32> + %expected1 = check.generate.fill value(36.0) : tensor<32xf32> + check.expect.equal actual(%row1) expected(%expected1) : tensor<32xf32> + %row2 = check.tensor.view %separate offset(256) : tensor<160xf32> -> tensor<32xf32> + %expected2 = check.generate.fill value(39.0) : tensor<32xf32> + check.expect.equal actual(%row2) expected(%expected2) : tensor<32xf32> + %row3 = check.tensor.view %separate offset(384) : tensor<160xf32> -> tensor<32xf32> + %expected3 = check.generate.fill value(40.0) : tensor<32xf32> + check.expect.equal actual(%row3) expected(%expected3) : tensor<32xf32> + %row4 = check.tensor.view %separate offset(512) : tensor<160xf32> -> tensor<32xf32> + %expected4 = check.generate.fill value(40.0) : tensor<32xf32> + check.expect.equal actual(%row4) expected(%expected4) : tensor<32xf32> + check.return +} 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 new file mode 100644 index 000000000000..db53e44317ff --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_f32.loom @@ -0,0 +1,26 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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)] + +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 + %output_width = index.add %lhs_width, %rhs_width : 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(%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/concat_dim0_strided_source_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_strided_source_f32.loom new file mode 100644 index 000000000000..410f527839b3 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/concat_dim0_strided_source_f32.loom @@ -0,0 +1,37 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.concat_dim0.strided_source_f32_body(%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_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)] +config.decl @ggml.concat_dim0_strided_source_f32.row_count : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.concat_dim0_strided_source_f32.row_ne1 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.concat_dim0_strided_source_f32.row_ne2 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.lhs_stride3 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.concat_dim0_strided_source_f32.rhs_stride1 : %value: index where [range(%value, 1, 1073741824)] +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)] + +kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_concat_dim0_strided_source_f32") @ggml_concat_dim0_strided_source_f32(%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_strided_source_f32.lhs_width : index + %rhs_width = config.get @ggml.concat_dim0_strided_source_f32.rhs_width : index + %row_count = config.get @ggml.concat_dim0_strided_source_f32.row_count : index + %output_width = index.add %lhs_width, %rhs_width : index + %element_count = index.mul %output_width, %row_count : 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_span: index, %rhs_span: index, %lhs: buffer, %rhs: buffer, %output: buffer) { + template.apply<@ggml.copy.concat_dim0.strided_source_f32_body>(%lhs_span, %rhs_span, %lhs, %rhs, %output) : (index, index, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f16_k16_major.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f16_k16_major.loom new file mode 100644 index 000000000000..a88b6d2af63a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f16_k16_major.loom @@ -0,0 +1,26 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.layout_k16_major.body(%token_count: index, %input: buffer, %output: buffer) + +// Exact K16-major F16 layout for native matrix fragments. The ordinary F16 +// alternate remains unchanged; this copy is a private dispatch transient. +config.decl @ggml.copy_f16_k16_major.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] +config.decl @ggml.copy_f16_k16_major.token_count : %value: index where [range(%value, 512, 2048), mul(%value, 512)] +config.decl @ggml.copy_f16_k16_major.input_is_f16 : %value: index where [range(%value, 0, 1)] +amdgpu.target @ggml_copy_f16_k16_major_target {subgroup_size = 32} + +kernel.def target(@ggml_copy_f16_k16_major_target) export("ggml_copy_f16_k16_major") @ggml_copy_f16_k16_major(%token_count: index) { + %one = index.constant 1 : index + %threads = index.constant 256 : index + %tile = index.constant 2048 : index + %rounding = index.constant 2047 : index + %width = config.get @ggml.copy_f16_k16_major.input_size : index + %elements = index.mul %token_count, %width : index + %rounded = index.add %elements, %rounding : index + %groups = index.div %rounded, %tile : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%token_count: index, %input: buffer, %output: buffer) { + template.apply<@ggml.copy.layout_k16_major.body>(%token_count, %input, %output) : (index, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f32.loom new file mode 100644 index 000000000000..3348ba12eca6 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f32.loom @@ -0,0 +1,30 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.vector_load_store.f32_body(%element_count: index, %source: buffer, %output: buffer) + +amdgpu.target @ggml_copy_f32_gfx11_wave64 {subgroup_size = 64} + +kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_f32") @ggml_copy_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %source: buffer, %output: buffer) { + template.apply<@ggml.copy.vector_load_store.f32_body>(%element_count, %source, %output) : (index, buffer, buffer) + kernel.return +} + +check.case public @ggml_copy_f32_small_case { + %four = check.literal value(4) : index + %source = check.generate.iota offset(-2.0) step(1.0) : tensor<4xf32> + %output = check.generate.fill value(0.0) : tensor<4xf32> + %expected = check.generate.iota offset(-2.0) step(1.0) : tensor<4xf32> + kernel.launch @ggml_copy_f32[%four](%four, %source, %output) : [index](index, tensor<4xf32>, tensor<4xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.return +} + +check.benchmark<@ggml_copy_f32_small_case> @ggml_copy_f32_small diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f32_f16.loom new file mode 100644 index 000000000000..a5c940a11c2f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/copy_f32_f16.loom @@ -0,0 +1,30 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.convert_f32_f16.body(%element_count: index, %source: buffer, %output: buffer) + +amdgpu.target @ggml_copy_f32_gfx11_wave64 {subgroup_size = 64} + +// Aligned tensor conversions used by existing F16-consuming matmul entries. +kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_f32_f16") @ggml_copy_f32_f16(%element_count: index) { + %one = index.constant 1 : index + %threads = index.constant 256 : index + %tile = index.constant 1024 : index + %rounding = index.constant 1023 : index + %rounded = index.add %element_count, %rounding : index + %groups = index.div %rounded, %tile : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%element_count: index, %source: buffer, %output: buffer) { + template.apply<@ggml.copy.convert_f32_f16.body>(%element_count, %source, %output) : (index, buffer, buffer) + kernel.return +} + +check.case public @ggml_copy_f32_f16_small_case { + %count = check.literal value(12) : index + %source = check.generate.iota offset(-2.0) step(0.5) : tensor<12xf32> + %output = check.generate.fill value(0.0) : tensor<12xf16> + %expected = check.generate.iota offset(-2.0) step(0.5) : tensor<12xf16> + kernel.launch @ggml_copy_f32_f16[%count](%count, %source, %output) : [index](index, tensor<12xf32>, tensor<12xf16>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<12xf16> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/strided_source_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/strided_source_f32.loom new file mode 100644 index 000000000000..00572d89307c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/strided_source_f32.loom @@ -0,0 +1,39 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.copy.strided_source.f32_body(%element_count: index, %source_span: index, %source: buffer, %output: buffer) + +amdgpu.target @ggml_copy_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.copy_strided_source_f32.ne0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.ne1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.ne2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride0 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride1 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride2 : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.copy_strided_source_f32.stride3 : %value: index where [range(%value, 1, 1073741824)] + +kernel.def target(@ggml_copy_f32_gfx11_wave64) export("ggml_copy_strided_source_f32") @ggml_copy_strided_source_f32(%element_count: index, %source_span: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %source_span: index, %source: buffer, %output: buffer) { + template.apply<@ggml.copy.strided_source.f32_body>(%element_count, %source_span, %source, %output) : (index, index, buffer, buffer) + kernel.return +} + +check.case public @ggml_copy_strided_source_f32_small_case { + %six = check.literal value(6) : index + %eight = check.literal value(8) : index + %source = check.generate.iota offset(0.0) step(1.0) period(5) : tensor<10xf32> + %output = check.generate.fill value(-1.0) : tensor<6xf32> + %expected = check.generate.iota offset(0.0) step(1.0) period(3) : tensor<6xf32> + kernel.launch @ggml_copy_strided_source_f32[%six, %eight](%six, %eight, %source, %output) : [index, index](index, index, tensor<10xf32>, tensor<6xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<6xf32> + check.return +} + +check.benchmark<@ggml_copy_strided_source_f32_small_case> @ggml_copy_strided_source_f32_small 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 new file mode 100644 index 000000000000..3b24f8946c4a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/copy/transpose_f16.loom @@ -0,0 +1,22 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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} + +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(%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_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_bc_f32.loom new file mode 100644 index 000000000000..ea3b63059209 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_bc_f32.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.elementwise.broadcast_index(%linear: index, %n0: index, %n1: index, %n2: index, %coord0: index, %coord1: index, %coord2: index, %coord3: index, %broadcast_dim0_flag: index, %broadcast_dim1_flag: index, %broadcast_dim2_flag: index, %broadcast_dim3_flag: index, %zero: index, %one: index) -> (index) + +amdgpu.target @ggml_binary_bc_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.binary_bc_f32.op : %value: index where [range(%value, 0, 8)] + +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.src0_broadcast_dim2 : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.binary_bc_f32.src0_broadcast_dim3 : %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)] + +config.decl @ggml.binary_bc_f32.src1_broadcast_dim2 : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.binary_bc_f32.src1_broadcast_dim3 : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_binary_bc_f32_gfx11_wave64) export("ggml_binary_bc_f32") @ggml_binary_bc_f32(%element_count: index, %ne0: index, %ne1: index, %ne2: index, %ne3: index, %src0_element_count: index, %src1_element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %ne0: index, %ne1: index, %ne2: index, %ne3: index, %src0_element_count: index, %src1_element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %n0 = index.assume %ne0 [range(%ne0, 1, 134217728)] : index + %n1 = index.assume %ne1 [range(%ne1, 1, 134217728)] : index + %n2 = index.assume %ne2 [range(%ne2, 1, 134217728)] : index + %src0_count = index.assume %src0_element_count [range(%src0_element_count, 1, 134217728)] : index + %src1_count = index.assume %src1_element_count [range(%src1_element_count, 1, 134217728)] : index + %op = config.get @ggml.binary_bc_f32.op : index + %src0_broadcast_dim0_flag = config.get @ggml.binary_bc_f32.src0_broadcast_dim0 : index + %src0_broadcast_dim1_flag = config.get @ggml.binary_bc_f32.src0_broadcast_dim1 : index + %src0_broadcast_dim2_flag = config.get @ggml.binary_bc_f32.src0_broadcast_dim2 : index + %src0_broadcast_dim3_flag = config.get @ggml.binary_bc_f32.src0_broadcast_dim3 : index + %src1_broadcast_dim0_flag = config.get @ggml.binary_bc_f32.src1_broadcast_dim0 : index + %src1_broadcast_dim1_flag = config.get @ggml.binary_bc_f32.src1_broadcast_dim1 : index + %src1_broadcast_dim2_flag = config.get @ggml.binary_bc_f32.src1_broadcast_dim2 : index + %src1_broadcast_dim3_flag = config.get @ggml.binary_bc_f32.src1_broadcast_dim3 : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %one = index.constant 1 : index + %zero = index.constant 0 : index + %base0 = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %lhs_view = buffer.view %lhs[%zero_offset] : buffer -> view<[%src0_count]xf32> + %rhs_view = buffer.view %rhs[%zero_offset] : buffer -> view<[%src1_count]xf32> + %output_view = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %coord0 = index.rem %linear, %n0 : index + %linear_div_n0 = index.div %linear, %n0 : index + %coord1 = index.rem %linear_div_n0, %n1 : index + %linear_div_n01 = index.div %linear_div_n0, %n1 : index + %coord2 = index.rem %linear_div_n01, %n2 : index + %coord3 = index.div %linear_div_n01, %n2 : index + + %src0_linear = template.apply<@ggml.elementwise.broadcast_index>(%linear, %n0, %n1, %n2, %coord0, %coord1, %coord2, %coord3, %src0_broadcast_dim0_flag, %src0_broadcast_dim1_flag, %src0_broadcast_dim2_flag, %src0_broadcast_dim3_flag, %zero, %one) : (index, index, index, index, index, index, index, index, index, index, index, index, index, index) -> (index) + %src1_linear = template.apply<@ggml.elementwise.broadcast_index>(%linear, %n0, %n1, %n2, %coord0, %coord1, %coord2, %coord3, %src1_broadcast_dim0_flag, %src1_broadcast_dim1_flag, %src1_broadcast_dim2_flag, %src1_broadcast_dim3_flag, %zero, %one) : (index, index, index, index, index, index, index, index, index, index, index, index, index, index) -> (index) + + %lhs_value = view.load %lhs_view[%src0_linear] : view<[%src0_count]xf32> -> f32 + %rhs_value = view.load %rhs_view[%src1_linear] : view<[%src1_count]xf32> -> f32 + %result = template.apply<@ggml.binary_f32.apply>(%op, %lhs_value, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/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_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_f32.loom new file mode 100644 index 000000000000..f8e74b98b4bd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_f32.loom @@ -0,0 +1,96 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.binary_f32.apply(%arg0: index, %arg1: f32, %arg2: f32) -> (f32) + +amdgpu.target @ggml_binary_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.binary_f32.op : %value: index where [range(%value, 0, 8)] + +config.decl @ggml.binary_f32.ne0 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.ne1 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.ne2 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src0_stride1 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src0_stride2 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src0_stride3 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src1_stride1 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src1_stride2 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src1_stride3 : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src0_span : %value: index where [range(%value, 1, 134217728)] + +config.decl @ggml.binary_f32.src1_span : %value: index where [range(%value, 1, 134217728)] + +kernel.def target(@ggml_binary_f32_gfx11_wave64) export("ggml_binary_f32") @ggml_binary_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %n0 = config.get @ggml.binary_f32.ne0 : index + %n1 = config.get @ggml.binary_f32.ne1 : index + %n2 = config.get @ggml.binary_f32.ne2 : index + %s0_1 = config.get @ggml.binary_f32.src0_stride1 : index + %s0_2 = config.get @ggml.binary_f32.src0_stride2 : index + %s0_3 = config.get @ggml.binary_f32.src0_stride3 : index + %s1_1 = config.get @ggml.binary_f32.src1_stride1 : index + %s1_2 = config.get @ggml.binary_f32.src1_stride2 : index + %s1_3 = config.get @ggml.binary_f32.src1_stride3 : index + %src0_count = config.get @ggml.binary_f32.src0_span : index + %src1_count = config.get @ggml.binary_f32.src1_span : index + %op = config.get @ggml.binary_f32.op : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %zero = index.constant 0 : index + %base0 = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %lhs_view = buffer.view %lhs[%zero_offset] : buffer -> view<[%src0_count]xf32> + %rhs_view = buffer.view %rhs[%zero_offset] : buffer -> view<[%src1_count]xf32> + %output_view = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %coord0 = index.rem %linear, %n0 : index + %linear_div_n0 = index.div %linear, %n0 : index + %coord1 = index.rem %linear_div_n0, %n1 : index + %linear_div_n01 = index.div %linear_div_n0, %n1 : index + %coord2 = index.rem %linear_div_n01, %n2 : index + %coord3 = index.div %linear_div_n01, %n2 : index + %src0_term1 = index.mul %coord1, %s0_1 : index + %src0_term2 = index.mul %coord2, %s0_2 : index + %src0_term3 = index.mul %coord3, %s0_3 : index + %src0_linear01 = index.add %coord0, %src0_term1 : index + %src0_linear012 = index.add %src0_linear01, %src0_term2 : index + %src0_linear = index.add %src0_linear012, %src0_term3 : index + %src1_term1 = index.mul %coord1, %s1_1 : index + %src1_term2 = index.mul %coord2, %s1_2 : index + %src1_term3 = index.mul %coord3, %s1_3 : index + %src1_linear01 = index.add %coord0, %src1_term1 : index + %src1_linear012 = index.add %src1_linear01, %src1_term2 : index + %src1_linear = index.add %src1_linear012, %src1_term3 : index + %src0_in_bounds = index.cmp ult, %src0_linear, %src0_count : index + scf.if %src0_in_bounds { + %src1_in_bounds = index.cmp ult, %src1_linear, %src1_count : index + scf.if %src1_in_bounds { + %lhs_value = view.load %lhs_view[%src0_linear] : view<[%src0_count]xf32> -> f32 + %rhs_value = view.load %rhs_view[%src1_linear] : view<[%src1_count]xf32> -> f32 + %result = template.apply<@ggml.binary_f32.apply>(%op, %lhs_value, %rhs_value) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/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/elementwise/binary_swiglu_symmetric_i4_k32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_swiglu_symmetric_i4_k32.loom new file mode 100644 index 000000000000..52012948e425 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/binary_swiglu_symmetric_i4_k32.loom @@ -0,0 +1,28 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.binary_swiglu_symmetric_i4_k32.body(%lhs: buffer, %rhs: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +amdgpu.target @ggml_binary_swiglu_i4_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.binary_swiglu_symmetric_i4.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 64)] + +config.decl @ggml.binary_swiglu_symmetric_i4.token_count : %value: index where [range(%value, 1, 16)] + +kernel.def target(@ggml_binary_swiglu_i4_gfx11_wave32) export("ggml_binary_swiglu_symmetric_i4_k32") @ggml_binary_swiglu_symmetric_i4_k32() { + %c1 = index.constant 1 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %input_size = config.get @ggml.binary_swiglu_symmetric_i4.input_size : index + %token_count = config.get @ggml.binary_swiglu_symmetric_i4.token_count : index + %element_count = index.mul %input_size, %token_count : index + %group_count = index.div %element_count, %c64 : index + %rounded_group_count = index.add %group_count, %c7 : index + %workgroup_count = index.div %rounded_group_count, %c8 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%lhs: buffer, %rhs: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) { + template.apply<@ggml.binary_swiglu_symmetric_i4_k32.body>(%lhs, %rhs, %output, %i4_qs, %i4_ds, %i4_sums) : (buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_add_f32.loom new file mode 100644 index 000000000000..dfd639bac97a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_add_f32.loom @@ -0,0 +1,96 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.elementwise.scale_apply(%value: f32, %scale: f32, %bias: f32) -> (f32) +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_scale_add_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.scale_add_f32.ne0 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.ne1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.ne2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_stride1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_stride2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_stride3 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.input_span : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_stride1 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_stride2 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_stride3 : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.residual_span : %value: index where [range(%value, 1, 134217728)] +config.decl @ggml.scale_add_f32.scale : f32 +config.decl @ggml.scale_add_f32.bias : f32 +config.decl @ggml.scale_add_f32.binary_op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.scale_add_f32.scaled_lhs : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_scale_add_f32_gfx11_wave64) export("ggml_scale_add_f32") @ggml_scale_add_f32(%element_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c255 = index.constant 255 : index + %rounded = index.add %element_count, %c255 : index + %workgroup_count = index.div %rounded, %c256 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%element_count: index, %input: buffer, %residual: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %n0 = config.get @ggml.scale_add_f32.ne0 : index + %n1 = config.get @ggml.scale_add_f32.ne1 : index + %n2 = config.get @ggml.scale_add_f32.ne2 : index + %input_stride1 = config.get @ggml.scale_add_f32.input_stride1 : index + %input_stride2 = config.get @ggml.scale_add_f32.input_stride2 : index + %input_stride3 = config.get @ggml.scale_add_f32.input_stride3 : index + %input_count = config.get @ggml.scale_add_f32.input_span : index + %residual_stride1 = config.get @ggml.scale_add_f32.residual_stride1 : index + %residual_stride2 = config.get @ggml.scale_add_f32.residual_stride2 : index + %residual_stride3 = config.get @ggml.scale_add_f32.residual_stride3 : index + %residual_count = config.get @ggml.scale_add_f32.residual_span : index + %scale = config.get @ggml.scale_add_f32.scale : f32 + %bias = config.get @ggml.scale_add_f32.bias : f32 + %binary_op = config.get @ggml.scale_add_f32.binary_op : index + %scaled_lhs = config.get @ggml.scale_add_f32.scaled_lhs : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c256 = index.constant 256 : index + %base0 = index.mul %workgroup, %c256 : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %c0_offset = index.constant 0 : offset + %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%input_count]xf32> + %residual_view = buffer.view %residual[%c0_offset] : buffer -> view<[%residual_count]xf32> + %output_view = buffer.view %output[%c0_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %coord0 = index.rem %linear, %n0 : index + %linear_div_n0 = index.div %linear, %n0 : index + %coord1 = index.rem %linear_div_n0, %n1 : index + %linear_div_n01 = index.div %linear_div_n0, %n1 : index + %coord2 = index.rem %linear_div_n01, %n2 : index + %coord3 = index.div %linear_div_n01, %n2 : index + %input_term1 = index.mul %coord1, %input_stride1 : index + %input_term2 = index.mul %coord2, %input_stride2 : index + %input_term3 = index.mul %coord3, %input_stride3 : index + %input_linear01 = index.add %coord0, %input_term1 : index + %input_linear012 = index.add %input_linear01, %input_term2 : index + %input_linear = index.add %input_linear012, %input_term3 : index + %residual_term1 = index.mul %coord1, %residual_stride1 : index + %residual_term2 = index.mul %coord2, %residual_stride2 : index + %residual_term3 = index.mul %coord3, %residual_stride3 : index + %residual_linear01 = index.add %coord0, %residual_term1 : index + %residual_linear012 = index.add %residual_linear01, %residual_term2 : index + %residual_linear = index.add %residual_linear012, %residual_term3 : index + %input_in_bounds = index.cmp ult, %input_linear, %input_count : index + scf.if %input_in_bounds { + %residual_in_bounds = index.cmp ult, %residual_linear, %residual_count : index + scf.if %residual_in_bounds { + %input_value = view.load %input_view[%input_linear] : view<[%input_count]xf32> -> f32 + %scaled = template.apply<@ggml.elementwise.scale_apply>(%input_value, %scale, %bias) : (f32, f32, f32) -> (f32) + %residual_value = view.load %residual_view[%residual_linear] : view<[%residual_count]xf32> -> f32 + %c1 = index.constant 1 : index + %is_scaled_lhs = index.cmp eq, %scaled_lhs, %c1 : index + %lhs = scf.select %is_scaled_lhs, %scaled, %residual_value : f32 + %rhs = scf.select %is_scaled_lhs, %residual_value, %scaled : f32 + %result = template.apply<@ggml.binary_f32.apply>(%binary_op, %lhs, %rhs) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_bias_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_bias_f32.loom new file mode 100644 index 000000000000..d015d58ce614 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_bias_f32.loom @@ -0,0 +1,51 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.elementwise.scale_apply(%value: f32, %scale: f32, %bias: f32) -> (f32) + +amdgpu.target @ggml_scale_bias_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.scale.scale : f32 + +config.decl @ggml.scale.bias : f32 + +kernel.def target(@ggml_scale_bias_f32_gfx11_wave64) export("ggml_scale_bias_f32") @ggml_scale_bias_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %input: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %scale = config.get @ggml.scale.scale : f32 + %bias = config.get @ggml.scale.bias : f32 + %zero_offset = index.constant 0 : offset + %input_view = buffer.view %input[%zero_offset] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %input_value = view.load %input_view[%linear] : view<[%count]xf32> -> f32 + %result = template.apply<@ggml.elementwise.scale_apply>(%input_value, %scale, %bias) : (f32, f32, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} + +check.case public @ggml_scale_bias_f32_small_case { + %four = check.literal value(4) : index + %input = check.generate.iota offset(-2.0) step(1.0) : tensor<4xf32> + %output = check.generate.fill value(0.0) : tensor<4xf32> + %expected = check.generate.iota offset(-3.0) step(2.0) : tensor<4xf32> + kernel.launch @ggml_scale_bias_f32[%four](%four, %input, %output) : [index](index, tensor<4xf32>, tensor<4xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<4xf32> + check.return +} + +check.benchmark<@ggml_scale_bias_f32_small_case> @ggml_scale_bias_f32_small diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_f32.loom new file mode 100644 index 000000000000..442708eb48c2 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/scale_f32.loom @@ -0,0 +1,40 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.elementwise.scale_apply(%value: f32, %scale: f32, %bias: f32) -> (f32) + +amdgpu.target @ggml_scale_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.scale_f32.scale : f32 + +config.decl @ggml.scale_f32.bias : f32 + +kernel.def target(@ggml_scale_f32_gfx11_wave64) export("ggml_scale_f32") @ggml_scale_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %input: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %scale = config.get @ggml.scale_f32.scale : f32 + %bias = config.get @ggml.scale_f32.bias : f32 + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base0 = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%zero_offset] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %value = view.load %input_view[%linear] : view<[%count]xf32> -> f32 + %result = template.apply<@ggml.elementwise.scale_apply>(%value, %scale, %bias) : (f32, f32, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom new file mode 100644 index 000000000000..b858a1650c06 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/elementwise/unary_f32.loom @@ -0,0 +1,37 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.unary_f32.apply(%arg0: index, %arg1: f32) -> (f32) + +amdgpu.target @ggml_unary_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.unary_f32.op : %value: index where [range(%value, 0, 23)] + +kernel.def target(@ggml_unary_f32_gfx11_wave64) export("ggml_unary_f32") @ggml_unary_f32(%element_count: index) { + %one = index.constant 1 : index + %twofiftysix = index.constant 256 : index + %rounding = index.constant 255 : index + %rounded = index.add %element_count, %rounding : index + %workgroup_count = index.div %rounded, %twofiftysix : index + kernel.launch.config workgroups(%workgroup_count, %one, %one) workgroup_size(%twofiftysix, %one, %one) : index +} launch(%element_count: index, %input: buffer, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 134217728)] : index + %op = config.get @ggml.unary_f32.op : index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %twofiftysix = index.constant 256 : index + %base0 = index.mul %workgroup, %twofiftysix : index + %linear0 = index.add %base0, %workitem : index + %linear = index.assume %linear0 [range(%linear0, 0, 134217983)] : index + %in_bounds = index.cmp ult, %linear, %count : index + %zero_offset = index.constant 0 : offset + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%zero_offset] : buffer -> view<[%count]xf32> + %output_view = buffer.view %output_noalias[%zero_offset] : buffer -> view<[%count]xf32> + scf.if %in_bounds { + %value = view.load %input_view[%linear] : view<[%count]xf32> -> f32 + %result = template.apply<@ggml.unary_f32.apply>(%op, %value) : (index, f32) -> (f32) + view.store %result, %output_view[%linear] : f32, view<[%count]xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom new file mode 100644 index 000000000000..d0b4461a26a7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/decode_split_f32_f16_wmma_next_q8.loom @@ -0,0 +1,381 @@ +// Qwen3 MoE grouped-query decode FlashAttention. +// +// Each workgroup processes one 64-token KV block for all GQA query heads that +// share a KV head. The workgroups publish online-softmax state, then the last +// arrival folds every block and resets the per-KV-head completion counter for +// the next invocation. Packing GQA heads removes redundant K/V traffic while +// split-K preserves enough parallelism for decode without another dispatch. +template.decl @ggml.flash_attention.decode_split.pack_completed_q8(%key_value_head: index, %output: buffer, %q8_output: buffer) + +template.decl @ggml.flash_attention.decode_split.produce_partials(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) + +template.decl @ggml.flash_attention.decode_split.produce_partials.active(%key_value_token_count: index, %partial_block_capacity0: index, %launched_block_count0: index, %query: buffer, %key: buffer, %value: buffer, %lane_mask: f32, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_completed.cooperative(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_completed.direct(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.reduce_fused(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) + +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_flash_attention_decode_split_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.qk_head_size : %value: index where [range(%value, 16, 512), mul(%value, 16)] + +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] + +config.decl @ggml.flash_attention.attention_scale : f32 + +// Maximum K/V storage capacity available to the compiled kernel. +config.decl @ggml.flash_attention.decode.key_value_token_capacity : %value: index where [range(%value, 64, 32768)] + +// Short-context export: produce and reduce in one dispatch. +kernel.def target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_f32_f16_wmma(%key_value_token_count: index) { + %key_value_head_count = config.get @ggml.flash_attention.key_value_head_count : index + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index + kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer) { + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 2048)] : index + %bounded_key_value_token_count, %launch_key_value_token_capacity = index.assume %key_value_token_count_in_range, %key_value_token_capacity [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index, index + %padded_key_value_token_capacity = index.add %launch_key_value_token_capacity, %c63 : index + %producer_block_count = index.div %padded_key_value_token_capacity, %c64 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@ggml.flash_attention.decode_split.produce_partials>(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.flash_attention.decode_split.reduce_fused>(%launch_key_value_token_capacity, %producer_block_count, %producer_block_count, %publish_q8, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %partial_output) : (index, index, index, i1, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Short-context export that publishes the F32 attention result and the Q8_1 +// representation consumed by the following output projection. +kernel.def target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count: index) { + %key_value_head_count = config.get @ggml.flash_attention.key_value_head_count : index + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index + kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) { + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 2048)] : index + %bounded_key_value_token_count, %launch_key_value_token_capacity = index.assume %key_value_token_count_in_range, %key_value_token_capacity [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index, index + %padded_key_value_token_capacity = index.add %launch_key_value_token_capacity, %c63 : index + %producer_block_count = index.div %padded_key_value_token_capacity, %c64 : index + %publish_q8 = scalar.constant true : i1 + template.apply<@ggml.flash_attention.decode_split.produce_partials>(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.flash_attention.decode_split.reduce_fused>(%launch_key_value_token_capacity, %producer_block_count, %producer_block_count, %publish_q8, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %next_q8_output) : (index, index, index, i1, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Long-context producer export: publish partials for a following parallel +// reducer without carrying short-context synchronization bindings. +kernel.def target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_produce_partials_f32_f16_wmma(%key_value_token_count: index) { + %key_value_head_count = config.get @ggml.flash_attention.key_value_head_count : index + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index + kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index + %bounded_key_value_token_count = index.assume %key_value_token_count_in_range [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index + template.apply<@ggml.flash_attention.decode_split.produce_partials>(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Long contexts have enough split-K partials that assigning the reduction to +// one last-arriving producer serializes useful work. This reducer launches one +// two-wave workgroup per query head after an execution barrier from the +// producer. The first wave computes normalization once and both waves consume +// it, avoiding the duplicate work of independent 64-channel output slices. +kernel.def target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_reduce_f32(%key_value_token_count: index) { + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %query_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%key_value_token_count: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) { + %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 + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %query_head0 = kernel.workgroup.id : index + %query_head = index.assume %query_head0 [range(%query_head0, 0, 63)] : index + %workitem = kernel.workitem.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c64 = index.constant 64 : index + %c0_offset = index.constant 0 : offset + %reduction_stage_bytes = index.constant 8 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -1e+30 : f32 + %key_value_token_capacity = config.get @ggml.flash_attention.decode.key_value_token_capacity : index + %c63 = index.constant 63 : index + %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index + %key_value_block_count0 = index.div %padded_key_value_token_capacity, %c64 : index + %key_value_block_count = index.assume %key_value_block_count0 [range(%key_value_block_count0, 1, 512)] : index + %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index + %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index + %key_value_head = index.div %query_head, %query_heads_per_key_value_head : index + %query_row = index.rem %query_head, %query_heads_per_key_value_head : index + %is_first_subgroup = index.cmp eq, %subgroup, %c0 : index + %workitem_is_zero = index.cmp eq, %workitem, %c0 : index + %output_channel = index.add %workitem, %c0 : index + %output_channel_valid = index.cmp ult, %output_channel, %value_head_size : index + %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %output : buffer, buffer, buffer, buffer + %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer + %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer + %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer + %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer + %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> + %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x[%value_head_size]xf32> + %reduction_stage = buffer.alloca align(8) %reduction_stage_bytes : buffer + %reduction_stage_view = buffer.view %reduction_stage[%c0_offset] : buffer -> view<2xf32> + %lane_maximum = scf.if %is_first_subgroup -> (f32) { + %maximum = scf.for %block = [%lane to %key_value_block_count step %c64](%running_maximum = %negative_large : f32) -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %next_maximum = scalar.maxnumf %running_maximum, %block_maximum : f32 + scf.yield %next_maximum : f32 + } + scf.yield %maximum : f32 + } else { + scf.yield %negative_large : f32 + } + %subgroup_maximum = kernel.subgroup.reduce %lane_maximum : f32 + scf.if %workitem_is_zero { + view.store %subgroup_maximum, %reduction_stage_view[%c0] : f32, view<2xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %maximum = view.load %reduction_stage_view[%c0] : view<2xf32> -> f32 + %lane_sum = scf.if %is_first_subgroup -> (f32) { + %sum = scf.for %block = [%lane to %key_value_block_count step %c64](%running_sum = %c0_f32 : f32) -> (f32) { + %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %delta = scalar.subf %block_maximum, %maximum : f32 + %scale = scalar.expf %delta : f32 + view.store %scale, %partial_max_view[%key_value_head, %block, %query_row] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> + %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %scaled_sum = scalar.mulf %block_sum, %scale : f32 + %next_sum = scalar.addf %running_sum, %scaled_sum : f32 + scf.yield %next_sum : f32 + } + scf.yield %sum : f32 + } else { + scf.yield %c0_f32 : f32 + } + %subgroup_sum = kernel.subgroup.reduce %lane_sum : f32 + scf.if %workitem_is_zero { + view.store %subgroup_sum, %reduction_stage_view[%c1] : f32, view<2xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %sum = view.load %reduction_stage_view[%c1] : view<2xf32> -> f32 + scf.if %output_channel_valid { + %unnormalized_output = scf.for %block = [%c0 to %key_value_block_count step %c1](%running_output = %c0_f32 : f32) -> (f32) unroll(%c4) schedule(interleaved) { + %scale = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 + %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%key_value_block_count]x16x[%value_head_size]xf16> -> f16 + %block_output = scalar.extf %block_output_f16 : f16 to f32 + %scaled_output = scalar.mulf %block_output, %scale : f32 + %next_output = scalar.addf %running_output, %scaled_output : f32 + scf.yield %next_output : f32 + } + %normalized_output = scalar.divf %unnormalized_output, %sum : f32 + view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x[%value_head_size]xf32> + } + kernel.return +} + +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_case { + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<256x4x128xf16> + %value = check.generate.fill value(2.0) : tensor<256x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<256xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x4x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(2.0) : tensor<32x128xf32> + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// Only KV row zero participates. The finite F32 iota step overflows to F16 +// negative infinity at every later row, leaving blocks one through three +// entirely masked. Each empty split must publish the online-softmax identity +// instead of evaluating -inf - -inf and contaminating the final reduction. +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_masked_blocks_case { + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<256x4x128xf16> + %value = check.generate.fill value(2.0) : tensor<256x4x128xf16> + %mask = check.generate.iota offset(0.0) step(-1e+30) : tensor<256xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x4x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x4x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(2.0) : tensor<32x128xf32> + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// The production Decode-513 shape combines eight full blocks with a one-row +// tail and selects the cooperative fused reducer. Constant nonzero V keeps the +// expected result exact while all four KV heads, 32 GQA heads, completion +// counters, partial tensors, and tail guards participate. A second invocation +// reuses the partial and counter storage with a different V tensor, making the +// completion-counter reset observable. +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_decode_513_case { + %key_value_token_count = check.literal value(513) : index + %query = check.generate.fill value(1.0) : tensor<32x128xf32> + %key = check.generate.fill value(1.0) : tensor<513x4x128xf16> + %value0 = check.generate.fill value(2.0) : tensor<513x4x128xf16> + %value1 = check.generate.fill value(3.0) : tensor<513x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<513xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output0 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %output1 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %expected0 = check.generate.fill value(2.0) : tensor<32x128xf32> + %expected1 = check.generate.fill value(3.0) : tensor<32x128xf32> + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output0) : [index](index, tensor<32x128xf32>, tensor<513x4x128xf16>, tensor<513x4x128xf16>, tensor<513xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output1) : [index](index, tensor<32x128xf32>, tensor<513x4x128xf16>, tensor<513x4x128xf16>, tensor<513xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output0) expected(%expected0) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.expect.close actual(%output1) expected(%expected1) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> + check.return +} + +// Sixty-five KV rows force a second split containing one valid row. The mask +// selects that final row, whose iota values begin at 1024, so an omitted or +// uninitialized tail cannot accidentally satisfy the check. +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_tail_case { + %key_value_token_count = check.literal value(65) : index + %query = check.generate.fill value(1.0) : tensor<1x128xf32> + %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> + %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<65xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<1x2x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<1x2x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<1x2x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<1xi32> + %output = check.generate.fill value(-1.0) : tensor<1x128xf32> + %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x128xf32> + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<65xf16>, tensor<1x2x16xf32>, tensor<1x2x16xf32>, tensor<1x2x16x128xf16>, tensor<1xi32>, tensor<1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> + check.return +} + +// Thirty-two split-K blocks exercise the separate parallel reducer and the +// execution barrier between its producer and consumer dispatches. +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_long_case { + %key_value_token_count = check.literal value(2048) : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<2048x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<2048x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<2048xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x32x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x32x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x32x16x128xf16> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @ggml_flash_attention_decode_split_produce_partials_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : [index](index, tensor<32x128xf32>, tensor<2048x4x128xf16>, tensor<2048x4x128xf16>, tensor<2048xf16>, tensor<4x32x16xf32>, tensor<4x32x16xf32>, tensor<4x32x16x128xf16>) + kernel.launch @ggml_flash_attention_decode_split_reduce_f32[%key_value_token_count](%key_value_token_count, %partial_max, %partial_sum, %partial_output, %output) : [index](index, tensor<4x32x16xf32>, tensor<4x32x16xf32>, tensor<4x32x16x128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case { + %key_value_token_count = check.param.choice values([64, 65, 128, 256, 512, 513, 768, 1024, 1280, 2048]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> + // Reserve the bounded scratch capacity once; each specialization addresses + // only ceildiv(key_value_token_count, 64) blocks. + %partial_max = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x512x16x128xf16> + %completion_counter = check.generate.fill value(0) : tensor<4xi32> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +// Long-context execution is one reusable producer/reducer command buffer. The +// harness records an explicit dispatch execution barrier between these calls +// and profiles their complete device-side span as one semantic operation. +check.case public @ggml_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case { + %key_value_token_count = check.param.choice values([2048, 32768]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> + %partial_max = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_sum = check.generate.fill value(-1.0) : tensor<4x512x16xf32> + %partial_output = check.generate.fill value(-1.0) : tensor<4x512x16x128xf16> + %output = check.generate.fill value(1.0) : tensor<32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<32x128xf32> + kernel.launch @ggml_flash_attention_decode_split_produce_partials_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>) + kernel.launch @ggml_flash_attention_decode_split_reduce_f32[%key_value_token_count](%key_value_token_count, %partial_max, %partial_sum, %partial_output, %output) : [index](index, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>, tensor<32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> + check.return +} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_64 {key_value_token_count = 64} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_65 {key_value_token_count = 65} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_128 {key_value_token_count = 128} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_256 {key_value_token_count = 256} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_masked_blocks_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_256_masked_blocks + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_512 {key_value_token_count = 512} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_513 {key_value_token_count = 513} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_768 {key_value_token_count = 768} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_1024 {key_value_token_count = 1024} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_1280 {key_value_token_count = 1280} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_2048_fused {key_value_token_count = 2048} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_2048 {key_value_token_count = 2048} + +check.benchmark<@ggml_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case> @ggml_flash_attention_decode_split_f32_f16_wmma_decode_32768 {key_value_token_count = 32768} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/decode_split_next_q8_test.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/decode_split_next_q8_test.loom new file mode 100644 index 000000000000..29be28a8919e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/decode_split_next_q8_test.loom @@ -0,0 +1,208 @@ +// Test-only differentials for fused decode-attention Q8 publication. The +// production kernels remain the actual providers; this module contributes only +// reference packing, masked metadata setup, and comparison cases. +template.decl @ggml.flash_attention.decode_split.quantize_reference_group(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) + +template.decl @ggml.flash_attention.decode_split.pack_completed_q8(%arg0: index, %arg1: buffer, %arg2: buffer) + +target.decl @ggml_flash_attention_decode_split_gfx11_wave64 + +kernel.decl target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_f32_f16_wmma(%key_value_token_count$8: index) launch(%key_value_token_count$9: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer) + +kernel.decl target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count$19: index) launch(%key_value_token_count$20: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) + +// Independent test-only implementation of one physical Q8_1 x4 group. This +// intentionally does not call the vector publication motif used in production. +template.def<@ggml.flash_attention.decode_split.quantize_reference_group> device @ggml_flash_attention_decode_split_quantize_reference_group(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { + %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index + %lane = kernel.workitem.id : 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 + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %group_bytes = index.constant 144 : offset + %payload_byte_add = index.constant 16 : offset + %scratch_d_byte_add = index.constant 128 : offset + %scratch_bytes = index.constant 144 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %c1_f32 = scalar.constant 1.0 : f32 + %c127 = scalar.constant 127.0 : f32 + %c0_offset = index.constant 0 : offset + %launched_element_count = index.mul %group_count, %c128 : index + %block_in_group0 = index.div %lane, %c8 : index + %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index + %word_in_block0 = index.rem %lane, %c8 : index + %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index + %group_element_base = index.mul %group, %c128 : index + %block_element_add = index.mul %block_in_group, %c32 : index + %block_word_add = index.mul %block_in_group, %c8 : index + %word_element_add = index.mul %word_in_block, %c4 : index + %input_block_base = index.add %group_element_base, %block_element_add : index + %input_index = index.add %input_block_base, %word_element_add : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> + %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> + %absolute_values = vector.absf %input_values : vector<4xf32> + %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> + view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index + scf.if %is_cohort_leader { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 + %d = scalar.divf %amax, %c127 : f32 + view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 + %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 + %d_inverse = scf.if %d_nonzero -> (f32) { + %inverse = scalar.divf %c1_f32, %d : f32 + scf.yield %inverse : f32 + } else { + scf.yield %c0_f32 : f32 + } + %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> + %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> + %rounded_values = vector.roundf %scaled_values : vector<4xf32> + %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> + %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> + %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset + %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset + %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> + %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> + %packed_word_index0 = index.add %block_word_add, %word_in_block : index + %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index + scf.if %publish_output { + vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> + } + %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 + view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 + scf.if %publishes_metadata { + %cohort_base = index.mul %block_in_group, %c8 : index + %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> + %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 + %s = scalar.mulf %quantized_sum, %d : f32 + %d_f16 = scalar.fptrunc %d : f32 to f16 + %s_f16 = scalar.fptrunc %s : f32 to f16 + %ds_index = index.mul %block_in_group, %c2 : index + view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> + %s_index = index.add %ds_index, %c1 : index + view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> + } + template.return +} + +// Publishes the canonical reference Q8_1 x4 layout for one 4096-element row. +kernel.def @ggml_flash_attention_decode_split_quantize_reference_4096() { + %c1 = index.constant 1 : index + %c32 = index.constant 32 : index + kernel.launch.config workgroups(%c32, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%input: buffer, %output: buffer) { + %publish_output = scalar.constant true : i1 + %group_count = index.constant 32 : index + %group = kernel.workgroup.id : index + template.apply<@ggml.flash_attention.decode_split.quantize_reference_group>(%publish_output, %group_count, %group, %input, %output) : (i1, index, index, buffer, buffer) + kernel.return +} + +// Isolates the publication phase for access-sanitized coverage without +// inflating the complete fused attention kernel beyond a short branch's range. +kernel.def target(@ggml_flash_attention_decode_split_gfx11_wave64) @ggml_flash_attention_decode_split_pack_completed_q8_test() { + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c4, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%input: buffer, %output: buffer) { + %key_value_head = kernel.workgroup.id : index + template.apply<@ggml.flash_attention.decode_split.pack_completed_q8>(%key_value_head, %input, %output) : (index, buffer, buffer) + kernel.return +} + +// Models the request-owned Decode-513 mask inside its 576-row capacity class. +kernel.def @ggml_flash_attention_decode_split_mask_513_of_576() { + %c1 = index.constant 1 : index + %c64 = index.constant 64 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%mask: buffer) { + %workitem = kernel.workitem.id : index + %c63 = index.constant 63 : index + %c513 = index.constant 513 : index + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %negative_large = scalar.constant -1e+30 : f32 + %negative_infinity = scalar.fptrunc %negative_large : f32 to f16 + %valid = index.cmp ult, %workitem, %c63 : index + %row0 = index.add %c513, %workitem : index + %row = scf.select %valid, %row0, %c0 : index + %mask_view = buffer.view %mask[%c0_offset] : buffer -> view<576xf16> + scf.if %valid { + view.store %negative_infinity, %mask_view[%row] : f16, view<576xf16> + } + kernel.return +} + +// The fused producer must match the ordinary attention export followed by the +// canonical GGML packer for every F32 value and every packed byte. Two value +// tensors reuse both paths' partials and completion counters, making counter +// reset and stale-publication failures observable. +check.case public @ggml_flash_attention_decode_split_next_q8_capacity_576_differential_case { + %key_value_token_count = check.literal value(576) : index + %query_seed = check.param.seed base(5858425849413783864) count(1) : i64 + %key_seed = check.param.seed base(5858425849313120568) count(1) : i64 + %value_seed0 = check.param.seed base(5858425849497669936) count(1) : i64 + %value_seed1 = check.param.seed base(5858425849497669937) count(1) : i64 + %query = check.generate.random.uniform seed(%query_seed) range(-1.0 to 1.0) : tensor<32x128xf32> + %key = check.generate.random.uniform seed(%key_seed) range(-1.0 to 1.0) : tensor<576x4x128xf16> + %value0 = check.generate.random.uniform seed(%value_seed0) range(-1.0 to 1.0) : tensor<576x4x128xf16> + %value1 = check.generate.random.uniform seed(%value_seed1) range(-1.0 to 1.0) : tensor<576x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<576xf16> + %reference_partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %reference_partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %reference_partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> + %reference_counter = check.generate.fill value(0) : tensor<4xi32> + %actual_partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %actual_partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> + %actual_partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> + %actual_counter = check.generate.fill value(0) : tensor<4xi32> + %reference_output0 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %reference_output1 = check.generate.fill value(-1.0) : tensor<32x128xf32> + %actual_output0 = check.generate.fill value(-2.0) : tensor<32x128xf32> + %actual_output1 = check.generate.fill value(-2.0) : tensor<32x128xf32> + %reference_q8_0 = check.generate.fill value(0) : tensor<4608xi8> + %reference_q8_1 = check.generate.fill value(0) : tensor<4608xi8> + %actual_q8_0 = check.generate.fill value(1) : tensor<4608xi8> + %actual_q8_1 = check.generate.fill value(1) : tensor<4608xi8> + kernel.launch @ggml_flash_attention_decode_split_mask_513_of_576(%mask) : (tensor<576xf16>) + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %reference_partial_max, %reference_partial_sum, %reference_partial_output, %reference_counter, %reference_output0) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + kernel.launch @ggml_flash_attention_decode_split_quantize_reference_4096(%reference_output0, %reference_q8_0) : (tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %actual_partial_max, %actual_partial_sum, %actual_partial_output, %actual_counter, %actual_output0, %actual_q8_0) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %reference_partial_max, %reference_partial_sum, %reference_partial_output, %reference_counter, %reference_output1) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) + kernel.launch @ggml_flash_attention_decode_split_quantize_reference_4096(%reference_output1, %reference_q8_1) : (tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @ggml_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %actual_partial_max, %actual_partial_sum, %actual_partial_output, %actual_counter, %actual_output1, %actual_q8_1) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) + check.expect.equal actual(%actual_output0) expected(%reference_output0) : tensor<32x128xf32> + check.expect.equal actual(%actual_q8_0) expected(%reference_q8_0) : tensor<4608xi8> + check.expect.equal actual(%actual_output1) expected(%reference_output1) : tensor<32x128xf32> + check.expect.equal actual(%actual_q8_1) expected(%reference_q8_1) : tensor<4608xi8> + check.return +} + +check.case public @ggml_flash_attention_decode_split_pack_completed_q8_differential_case { + %input_seed = check.param.seed base(5858425849396089144) count(1) : i64 + %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<32x128xf32> + %reference = check.generate.fill value(0) : tensor<4608xi8> + %actual = check.generate.fill value(1) : tensor<4608xi8> + kernel.launch @ggml_flash_attention_decode_split_quantize_reference_4096(%input, %reference) : (tensor<32x128xf32>, tensor<4608xi8>) + kernel.launch @ggml_flash_attention_decode_split_pack_completed_q8_test(%input, %actual) : (tensor<32x128xf32>, tensor<4608xi8>) + check.expect.equal actual(%actual) expected(%reference) : tensor<4608xi8> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom new file mode 100644 index 000000000000..d687e7089325 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma.loom @@ -0,0 +1,424 @@ +// ggml grouped-query FlashAttention for the F32-query/F16-cache path. +// +// One four-wave workgroup computes 16 query rows for one query head and one +// 256-channel output group against 64 KV rows at a time. The ownership changes +// mirror the cooperative-matrix schedule used by llama.cpp's Vulkan CM1 kernel: +// +// 1. All workitems stage a scaled 16xQK-head-size F16 query tile. +// 2. Each wave computes one 16x16 QK score slice. +// 3. Scores cross LDS so each wave can normalize four complete query rows. +// 4. F16 probabilities cross LDS for four P*V WMMA steps. +// 5. Each active lane retains one four-channel F16 packet for each of its +// four query rows across subsequent 64-row KV blocks. +// +// K and V remain in the row-major llama.cpp cache layout +// [KV token][KV head][head_size]. K uses the QK head size; V uses the output +// value head size. Their aligned F16 fragments load directly from global +// memory; there is no expanded or repacked persistent allocation. +// QK and the online-softmax statistics remain F32, while the P*V accumulation +// and carried output match the Vulkan oracle's F16 policy. +amdgpu.target @ggml_flash_attention_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.flash_attention.qk_head_size : %value: index where [range(%value, 16, 512), mul(%value, 16)] + +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] + +config.decl @ggml.flash_attention.attention_scale : f32 + +config.decl @ggml.flash_attention.apply_gate : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.flash_attention.gate_stride_head : %value: index where [range(%value, 1, 16777216)] + +config.decl @ggml.flash_attention.gate_stride_token : %value: index where [range(%value, 1, 16777216)] + +// Bounds the context copied by the test-only row-extraction kernel below. +config.decl @ggml.flash_attention.test.context_capacity : %value: index where [range(%value, 1, 32768)] + +template.decl @ggml.flash_attention.prefill.stage_query(%query_token_count: index, %query_origin: index, %query_head: index, %query_load_iteration_count: index, %query_stage_stride: index, %qk_head_size: index, %attention_scale: f32, %query: buffer, %query_stage: buffer) + +template.decl @ggml.flash_attention.qk_block_wmma.prefill_full(%bounded_key_value_token_count: index, %key_width: index, %qk_head_size: index, %query_stage_stride: index, %score_key_origin: index, %key_head_base: index, %subgroup_score_column: index, %key: buffer, %query_stage: buffer, %score_stage: buffer) + +template.decl @ggml.flash_attention.qk_block_wmma.prefill_tail(%qk_head_size: index, %query_stage_stride: index, %tail_stage_head_size: index, %tail_score_wave: i1, %subgroup: index, %tail_key_value_stage: buffer, %query_stage: buffer, %score_stage: buffer) + +template.decl @ggml.flash_attention.online_softmax.prefill(%masked_scores: vector<4xf32>, %valid: vector<4xi1>, %current_max: vector<4xf32>, %current_sum: vector<4xf32>, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %lane: index, %probability_stage: buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>) + +template.decl @ggml.flash_attention.pv_block_wmma.prefill_full(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %output_group_tile_base: index, %output_tile_count: index, %value_head_base: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %key_origin: index, %last_full_key_tile_start: index, %bounded_key_value_token_count: index, %value_width: index, %value_row_stride: index, %value_column_stride: index, %value: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + +template.decl @ggml.flash_attention.pv_block_wmma.prefill_tail(%old_scale: vector<4xf32>, %current_output0: vector<4xf16>, %current_output1: vector<4xf16>, %current_output2: vector<4xf16>, %current_output3: vector<4xf16>, %group_output_tile_count: index, %subgroup_product_channel: index, %lane_output_tile: index, %lane_product_channel: index, %query_row0: index, %query_row1: index, %query_row2: index, %query_row3: index, %tail_stage_head_size: index, %tail_key_value_stage: buffer, %probability_stage: buffer, %product_stage: buffer) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) + +template.decl @ggml.flash_attention.prefill.publish_output(%publish_f16: i1, %query_token_count: index, %query_head_count: index, %value_head_size: index, %lane_has_output: i1, %query_valid0: i1, %query_valid1: i1, %query_valid2: i1, %query_valid3: i1, %query_token0: index, %query_token1: index, %query_token2: index, %query_token3: index, %query_head: index, %lane_output_channel: index, %final_sum: vector<4xf32>, %final_output0: vector<4xf16>, %final_output1: vector<4xf16>, %final_output2: vector<4xf16>, %final_output3: vector<4xf16>, %gate: buffer, %output: buffer, %f16_output: buffer) + +template.decl @ggml.flash_attention.prefill.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) + +kernel.def target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) { + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_query_token_count = index.add %query_token_count, %c15 : index + %query_tile_count = index.div %padded_query_token_count, %c16 : index + %output_tile_count = index.div %value_head_size, %c64 : index + %padded_output_tile_count = index.add %output_tile_count, %c3 : index + %output_group_count = index.div %padded_output_tile_count, %c4 : index + kernel.launch.config workgroups(%query_tile_count, %query_head_count, %output_group_count) workgroup_size(%c256, %c1, %c1) : index +} launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) where [range(%query_token_count, 1, 2048)] { + %publish_f16 = scalar.constant false : i1 + template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output, %output) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Test-only causal-mask construction for the production 14-token witness. +// Finite F16 minima retain exact zero probabilities without requiring infinity +// support from synthetic tensor generators. Selective linking drops this +// helper from production roots. +kernel.def target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_test_make_causal_mask() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%mask: buffer) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c14 = index.constant 14 : index + %c196 = index.constant 196 : index + %c0_offset = index.constant 0 : offset + %c0_f16 = scalar.constant 0.0 : f16 + %negative_f16 = scalar.constant -65504.0 : f16 + %workitem = kernel.workitem.id : index + %in_bounds = index.cmp ult, %workitem, %c196 : index + %mask_noalias = buffer.assume.noalias %mask : buffer + %mask_view = buffer.view %mask_noalias[%c0_offset] : buffer -> view<14x14xf16> + scf.if %in_bounds { + %row = index.div %workitem, %c14 : index + %column = index.rem %workitem, %c14 : index + %row_limit = index.add %row, %c1 : index + %is_visible = index.cmp ult, %column, %row_limit : index + %value = scf.select %is_visible, %c0_f16, %negative_f16 : f16 + view.store %value, %mask_view[%row, %column] : f16, view<14x14xf16> + } + kernel.return +} + +// Test-only row extraction for comparing one multirow attention dispatch with +// independent one-row dispatches over identical data. Selective linking drops +// this helper from production roots. +kernel.def target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_test_extract_row(%query_token_count: index, %context_count: index, %source_row: index) { + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %qk_head_size0 = config.get @ggml.flash_attention.qk_head_size : index + %qk_head_size = index.assume %qk_head_size0 [range(%qk_head_size0, 16, 512), mul(%qk_head_size0, 16)] : index + %context_capacity = config.get @ggml.flash_attention.test.context_capacity : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c1 = index.constant 1 : index + %query_element_count = index.mul %query_head_count, %qk_head_size : index + %element_count = index.add %query_element_count, %context_capacity : index + %padded_element_count = index.add %element_count, %c255 : index + %workgroup_count = index.div %padded_element_count, %c256 : index + kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%query_token_count: index, %context_count: index, %source_row: index, %source_query: buffer, %source_mask: buffer, %target_query: buffer, %target_mask: buffer) { + %c0 = index.constant 0 : index + %c0_offset = index.constant 0 : offset + %c256 = index.constant 256 : index + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %qk_head_size0 = config.get @ggml.flash_attention.qk_head_size : index + %qk_head_size = index.assume %qk_head_size0 [range(%qk_head_size0, 16, 512), mul(%qk_head_size0, 16)] : index + %context_capacity = config.get @ggml.flash_attention.test.context_capacity : index + %bounded_query_token_count = index.assume %query_token_count [range(%query_token_count, 1, 2048)] : index + %bounded_context_count = index.assume %context_count [range(%context_count, 1, 32768), le(%context_count, %context_capacity)] : index + %bounded_source_row, %source_query_token_count = index.assume %source_row, %bounded_query_token_count [range(%source_row, 0, 2047), lt(%source_row, %bounded_query_token_count)] : index, index + %workgroup = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %linear = index.madd %workgroup, %c256, %workitem : index + %query_element_count = index.mul %query_head_count, %qk_head_size : index + %element_count = index.add %query_element_count, %bounded_context_count : index + %in_bounds = index.cmp ult, %linear, %element_count : index + %source_query_noalias, %source_mask_noalias, %target_query_noalias, %target_mask_noalias = buffer.assume.noalias %source_query, %source_mask, %target_query, %target_mask : buffer, buffer, buffer, buffer + %source_query_view = buffer.view %source_query_noalias[%c0_offset] : buffer -> view<[%source_query_token_count]x[%query_head_count]x[%qk_head_size]xf32> + %source_mask_view = buffer.view %source_mask_noalias[%c0_offset] : buffer -> view<[%source_query_token_count]x[%bounded_context_count]xf16> + %target_query_view = buffer.view %target_query_noalias[%c0_offset] : buffer -> view<1x[%query_head_count]x[%qk_head_size]xf32> + %target_mask_view = buffer.view %target_mask_noalias[%c0_offset] : buffer -> view<1x[%bounded_context_count]xf16> + scf.if %in_bounds { + %is_query_element = index.cmp ult, %linear, %query_element_count : index + scf.if %is_query_element { + %query_head = index.div %linear, %qk_head_size : index + %query_channel = index.rem %linear, %qk_head_size : index + %value = view.load %source_query_view[%bounded_source_row, %query_head, %query_channel] : view<[%source_query_token_count]x[%query_head_count]x[%qk_head_size]xf32> -> f32 + view.store %value, %target_query_view[%c0, %query_head, %query_channel] : f32, view<1x[%query_head_count]x[%qk_head_size]xf32> + } else { + %mask_column = index.sub %linear, %query_element_count : index + %value = view.load %source_mask_view[%bounded_source_row, %mask_column] : view<[%source_query_token_count]x[%bounded_context_count]xf16> -> f16 + view.store %value, %target_mask_view[%c0, %mask_column] : f16, view<1x[%bounded_context_count]xf16> + } + } + kernel.return +} + +// The mask selects the first KV row exactly. QK, F16 probability conversion, +// GQA addressing, and the P*V path all execute, while the expected result is +// the first V row and remains auditable as an iota. +check.case public @ggml_flash_attention_f32_f16_wmma_selected_row_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(64) : index + %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<64x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<1x64xf16> + %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x1x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<1x64xf16>, tensor<1x1x128xf32>, tensor<1x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> + check.return +} + +// Four query rows select four distinct KV rows. The 33-element mask period +// maps row-major index row*32+column to zero exactly when row == column for +// rows zero through three. The expected output is therefore the first four +// rows of V, preserving an auditable iota while distinguishing every component +// in the four-row per-wave ownership path. +check.case public @ggml_flash_attention_f32_f16_wmma_multirow_selected_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(32) : index + %query = check.generate.fill value(1.0) : tensor<4x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<32x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<32x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-10000.0) period(33) : tensor<4x32xf16> + %output = check.generate.fill value(-1.0) : tensor<4x1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<4x1x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<4x1x128xf32>, tensor<32x1x128xf16>, tensor<32x1x128xf16>, tensor<4x32xf16>, tensor<4x1x128xf32>, tensor<4x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<4x1x128xf32> + check.return +} + +// A unit mask step gives every row the same geometric softmax distribution, +// shifted to begin at the row-matched KV entry. Values advance by 1/4096, so +// each KV row adds exactly 1/32 and the infinite-series weighted row offset is +// 1/(32*(e-1)). Terms that wrap at period 33 are below F16 significance. This +// exercises four independent online-softmax and P*V states while retaining a +// compact closed-form expected iota. +check.case public @ggml_flash_attention_f32_f16_wmma_multirow_online_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(32) : index + %query = check.generate.fill value(1.0) : tensor<4x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<32x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.000244140625) : tensor<32x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-1.0) period(33) : tensor<4x32xf16> + %output = check.generate.fill value(-1.0) : tensor<4x1x128xf32> + %expected = check.generate.iota offset(0.01818677224124075) step(0.000244140625) : tensor<4x1x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<4x1x128xf32>, tensor<32x1x128xf16>, tensor<32x1x128xf16>, tensor<4x32xf16>, tensor<4x1x128xf32>, tensor<4x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<4x1x128xf32> + check.return +} + +// Compare the first four rows of one fourteen-row launch with four one-row +// launches over identical, nonuniform data and the production causal mask. +// Fourteen rows retain the exact dynamic dense-row stride, ownership, and +// query-validity pressure, while the one-row path is the proven containment. +// Row extraction only reshapes bindings and performs no attention arithmetic. +check.case public @ggml_flash_attention_f32_f16_wmma_multirow_differential_case { + %query_token_count = check.literal value(14) : index + %c1 = check.literal value(1) : index + %key_value_token_count = check.literal value(14) : index + %row0 = check.literal value(0) : index + %row1 = check.literal value(1) : index + %row2 = check.literal value(2) : index + %row3 = check.literal value(3) : index + %query_seed = check.param.seed base(5858425849414763858) count(1) : i64 + %key_seed = check.param.seed base(5858425849313057073) count(1) : i64 + %value_seed = check.param.seed base(5858425849497340977) count(1) : i64 + %query = check.generate.random.uniform seed(%query_seed) range(-1.0 to 1.0) : tensor<14x32x128xf32> + %key = check.generate.random.uniform seed(%key_seed) range(-1.0 to 1.0) : tensor<14x4x128xf16> + %value = check.generate.random.uniform seed(%value_seed) range(-1.0 to 1.0) : tensor<14x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<14x14xf16> + %actual = check.generate.fill value(-1.0) : tensor<14x32x128xf32> + %row_query = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %row_mask = check.generate.fill value(0.0) : tensor<1x14xf16> + %expected0 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected1 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected2 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %expected3 = check.generate.fill value(0.0) : tensor<1x32x128xf32> + %actual0 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + %actual1 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + %actual2 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + %actual3 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> + kernel.launch @ggml_flash_attention_test_make_causal_mask(%mask) : (tensor<14x14xf16>) + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %actual) : [index, index](index, index, tensor<14x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<14x14xf16>, tensor<14x32x128xf32>, tensor<14x32x128xf32>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row0](%query_token_count, %key_value_token_count, %row0, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %row_query, %expected0) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>, tensor<1x32x128xf32>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row0](%query_token_count, %key_value_token_count, %row0, %actual, %mask, %actual0, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row1](%query_token_count, %key_value_token_count, %row1, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %row_query, %expected1) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>, tensor<1x32x128xf32>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row1](%query_token_count, %key_value_token_count, %row1, %actual, %mask, %actual1, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row2](%query_token_count, %key_value_token_count, %row2, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %row_query, %expected2) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>, tensor<1x32x128xf32>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row2](%query_token_count, %key_value_token_count, %row2, %actual, %mask, %actual2, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row3](%query_token_count, %key_value_token_count, %row3, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + kernel.launch @ggml_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %row_query, %expected3) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>, tensor<1x32x128xf32>) + kernel.launch @ggml_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row3](%query_token_count, %key_value_token_count, %row3, %actual, %mask, %actual3, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) + check.expect.close actual(%actual0) expected(%expected0) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual1) expected(%expected1) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual2) expected(%expected2) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.expect.close actual(%actual3) expected(%expected3) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> + check.return +} + +// Seventeen rows force a partial second query tile. Equal scores and constant +// values make every valid output exactly two while still exercising online +// normalization and the query-tail guards. +check.case public @ggml_flash_attention_f32_f16_wmma_query_tail_case { + %query_token_count = check.literal value(17) : index + %key_value_token_count = check.literal value(64) : index + %query = check.generate.fill value(1.0) : tensor<17x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> + %value = check.generate.fill value(2.0) : tensor<64x1x128xf16> + %mask = check.generate.fill value(0.0) : tensor<17x64xf16> + %output = check.generate.fill value(-1.0) : tensor<17x1x128xf32> + %expected = check.generate.fill value(2.0) : tensor<17x1x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<17x1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<17x64xf16>, tensor<17x1x128xf32>, tensor<17x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<17x1x128xf32> + check.return +} + +// Sixty-five KV rows force one masked cleanup tile after a full WMMA block. +// The mask selects only that final row, whose iota values begin at 1024, so +// omitting the cleanup path cannot accidentally satisfy the check. +check.case public @ggml_flash_attention_f32_f16_wmma_key_value_tail_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(65) : index + %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> + %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<1x65xf16> + %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> + %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x1x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<1x65xf16>, tensor<1x1x128xf32>, tensor<1x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> + check.return +} + +// The first 64-row block contains the selected value while every mask in the +// second block rounds to negative infinity. This preserves a finite online +// softmax state while exercising the workgroup-wide block-pruning path. +check.case public @ggml_flash_attention_f32_f16_wmma_pruned_block_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(128) : index + %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> + %key = check.generate.fill value(1.0) : tensor<128x1x128xf16> + %value = check.generate.iota offset(0.0) step(0.125) : tensor<128x1x128xf16> + %mask = check.generate.iota offset(0.0) step(-2000.0) : tensor<1x128xf16> + %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> + %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x1x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<128x1x128xf16>, tensor<128x1x128xf16>, tensor<1x128xf16>, tensor<1x1x128xf32>, tensor<1x1x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> + check.return +} + +check.case public @ggml_flash_attention_f32_f16_wmma_benchmark_case { + %query_token_count = check.param.choice values([1, 32, 64, 128, 192, 255, 256, 257, 384, 511, 512, 513, 768, 1023, 1024, 1025, 1280, 1536, 1792, 2048]) name("query_token_count") : index + %key_value_token_count = check.param.choice values([64, 128, 192, 255, 256, 257, 384, 511, 512, 513, 768, 1023, 1024, 1025, 1280, 1536, 1792, 2048, 32768]) name("key_value_token_count") : index + %query = check.generate.fill value(0.0) : tensor<[%query_token_count]x32x128xf32> + %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> + %mask = check.generate.fill value(0.0) : tensor<[%query_token_count]x[%key_value_token_count]xf16> + %output = check.generate.fill value(1.0) : tensor<[%query_token_count]x32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<[%query_token_count]x32x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<[%query_token_count]x32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%query_token_count]x[%key_value_token_count]xf16>, tensor<[%query_token_count]x32x128xf32>, tensor<[%query_token_count]x32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%query_token_count]x32x128xf32> + check.return +} + +// Four finite 64-row blocks followed by four negative-infinity blocks model +// the aggregate computed/pruned work ratio of 512-token causal prefill while +// keeping every workgroup's path identical for a stable microbenchmark. +check.case public @ggml_flash_attention_f32_f16_wmma_pruned_half_benchmark_case { + %query_token_count = check.literal value(512) : index + %key_value_token_count = check.literal value(512) : index + %query = check.generate.fill value(0.0) : tensor<512x32x128xf32> + %key = check.generate.fill value(0.0) : tensor<512x4x128xf16> + %value = check.generate.fill value(0.0) : tensor<512x4x128xf16> + %mask = check.generate.iota offset(0.0) step(-300.0) period(512) : tensor<512x512xf16> + %output = check.generate.fill value(1.0) : tensor<512x32x128xf32> + %expected = check.generate.fill value(0.0) : tensor<512x32x128xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %query, %output) : [index, index](index, index, tensor<512x32x128xf32>, tensor<512x4x128xf16>, tensor<512x4x128xf16>, tensor<512x512xf16>, tensor<512x32x128xf32>, tensor<512x32x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<512x32x128xf32> + check.return +} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_selected_row_case> @ggml_flash_attention_f32_f16_wmma_selected_row + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_query_tail_case> @ggml_flash_attention_f32_f16_wmma_query_tail + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_key_value_tail_case> @ggml_flash_attention_f32_f16_wmma_key_value_tail + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_pruned_block_case> @ggml_flash_attention_f32_f16_wmma_pruned_block + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_pruned_half_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_512_pruned_half + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_decode_256 {key_value_token_count = 256, query_token_count = 1} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_decode_2048 {key_value_token_count = 2048, query_token_count = 1} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_decode_32768 {key_value_token_count = 32768, query_token_count = 1} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_32 {key_value_token_count = 256, query_token_count = 32} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_128 {key_value_token_count = 256, query_token_count = 128} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_512 {key_value_token_count = 512, query_token_count = 512} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_512_context_1024 {key_value_token_count = 1024, query_token_count = 512} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_512_context_1536 {key_value_token_count = 1536, query_token_count = 512} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_512_context_2048 {key_value_token_count = 2048, query_token_count = 512} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_1024 {key_value_token_count = 1024, query_token_count = 1024} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_prefill_2048 {key_value_token_count = 2048, query_token_count = 2048} + +// These aligned and boundary-adjacent self-attention shapes expose launch or +// tail cliffs that a powers-of-two-only benchmark would hide. They are +// measurement witnesses for one shape-specialized kernel, not routing buckets. +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_64 {key_value_token_count = 64, query_token_count = 64} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_128 {key_value_token_count = 128, query_token_count = 128} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_192 {key_value_token_count = 192, query_token_count = 192} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_255 {key_value_token_count = 255, query_token_count = 255} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_256 {key_value_token_count = 256, query_token_count = 256} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_257 {key_value_token_count = 257, query_token_count = 257} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_384 {key_value_token_count = 384, query_token_count = 384} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_511 {key_value_token_count = 511, query_token_count = 511} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_512 {key_value_token_count = 512, query_token_count = 512} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_513 {key_value_token_count = 513, query_token_count = 513} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_768 {key_value_token_count = 768, query_token_count = 768} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_1023 {key_value_token_count = 1023, query_token_count = 1023} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_1024 {key_value_token_count = 1024, query_token_count = 1024} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_1025 {key_value_token_count = 1025, query_token_count = 1025} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_1280 {key_value_token_count = 1280, query_token_count = 1280} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_1536 {key_value_token_count = 1536, query_token_count = 1536} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_1792 {key_value_token_count = 1792, query_token_count = 1792} + +check.benchmark<@ggml_flash_attention_f32_f16_wmma_benchmark_case> @ggml_flash_attention_f32_f16_wmma_self_2048 {key_value_token_count = 2048, query_token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom new file mode 100644 index 000000000000..409837c5b39f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/flash_attention/prefill_f32_f16_wmma_publish_f16.loom @@ -0,0 +1,38 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @ggml_flash_attention_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.flash_attention.query_head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.flash_attention.key_value_head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.flash_attention.qk_head_size : %value: index where [range(%value, 16, 512), mul(%value, 16)] +config.decl @ggml.flash_attention.value_head_size : %value: index where [range(%value, 64, 512), mul(%value, 64)] +config.decl @ggml.flash_attention.attention_scale : f32 +config.decl @ggml.flash_attention.apply_gate : %value: index where [range(%value, 0, 1)] +config.decl @ggml.flash_attention.gate_stride_head : %value: index where [range(%value, 1, 16777216)] +config.decl @ggml.flash_attention.gate_stride_token : %value: index where [range(%value, 1, 16777216)] + +template.decl @ggml.flash_attention.prefill.body(%publish_f16: i1, %query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) + +kernel.def target(@ggml_flash_attention_gfx11_wave64) export("ggml_flash_attention_f32_f16_wmma_publish_f16") @ggml_flash_attention_f32_f16_wmma_publish_f16(%query_token_count: index, %key_value_token_count: index) { + %query_head_count = config.get @ggml.flash_attention.query_head_count : index + %value_head_size0 = config.get @ggml.flash_attention.value_head_size : index + %value_head_size = index.assume %value_head_size0 [range(%value_head_size0, 64, 512), mul(%value_head_size0, 64)] : index + %c1 = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_query_token_count = index.add %query_token_count, %c15 : index + %query_tile_count = index.div %padded_query_token_count, %c16 : index + %output_tile_count = index.div %value_head_size, %c64 : index + %padded_output_tile_count = index.add %output_tile_count, %c3 : index + %output_group_count = index.div %padded_output_tile_count, %c4 : index + kernel.launch.config workgroups(%query_tile_count, %query_head_count, %output_group_count) workgroup_size(%c256, %c1, %c1) : index +} launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer, %f16_output: buffer) where [range(%query_token_count, 1, 2048)] { + %publish_f16 = scalar.constant true : i1 + template.apply<@ggml.flash_attention.prefill.body>(%publish_f16, %query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output, %f16_output) : (i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_inplace_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_inplace_f32_wmma.loom new file mode 100644 index 000000000000..29de4a4043fa --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_inplace_f32_wmma.loom @@ -0,0 +1,36 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +// Decode specialization keeps recurrent state in the graph-owned cache row. +kernel.def export("llm_gated_delta_net_f32_wmma_head128_inplace") @llm_gated_delta_net_f32_wmma_head128_inplace() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_inout: buffer, %dst: buffer) { + %state_inplace = scalar.constant true : i1 + %publish_snapshots = scalar.constant false : i1 + %projection_epilogue = scalar.constant false : i1 + %snapshot_stride = index.constant 1 : index + %scale = scalar.constant 0.0883883461356163 : f32 + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %g, %beta, %g, %beta, %state_inout, %dst, %state_inout, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_inplace_projection_epilogue_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_inplace_projection_epilogue_f32_wmma.loom new file mode 100644 index 000000000000..d9a4244cd60d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_inplace_projection_epilogue_f32_wmma.loom @@ -0,0 +1,37 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +// Low-row in-place scan with the alpha/beta projection epilogue folded into +// the existing one-thread-per-token scalar load point. +kernel.def export("llm_gated_delta_net_f32_wmma_head128_inplace_projection_epilogue") @llm_gated_delta_net_f32_wmma_head128_inplace_projection_epilogue() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %state_inout: buffer, %dst: buffer) { + %state_inplace = scalar.constant true : i1 + %publish_snapshots = scalar.constant false : i1 + %projection_epilogue = scalar.constant true : i1 + %snapshot_stride = index.constant 1 : index + %scale = scalar.constant 0.0883883461356163 : f32 + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %alpha_raw, %beta_raw, %bias, %a_scale, %state_inout, %dst, %state_inout, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_f32_wmma.loom new file mode 100644 index 000000000000..a717b7bfb9f6 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_f32_wmma.loom @@ -0,0 +1,35 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +kernel.def export("llm_gated_delta_net_f32_wmma_head128") @llm_gated_delta_net_f32_wmma_head128() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %dst: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant false : i1 + %projection_epilogue = scalar.constant false : i1 + %snapshot_stride = index.constant 1 : index + %scale = scalar.constant 0.0883883461356163 : f32 + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %g, %beta, %g, %beta, %state_in, %dst, %state_in, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_projection_epilogue_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_projection_epilogue_f32_wmma.loom new file mode 100644 index 000000000000..ea6bc3ec0ed0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_projection_epilogue_f32_wmma.loom @@ -0,0 +1,37 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +// Low-row scan with the alpha/beta projection epilogue folded into the +// existing one-thread-per-token scalar load point. +kernel.def export("llm_gated_delta_net_f32_wmma_head128_projection_epilogue") @llm_gated_delta_net_f32_wmma_head128_projection_epilogue() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant false : i1 + %projection_epilogue = scalar.constant true : i1 + %snapshot_stride = index.constant 1 : index + %scale = scalar.constant 0.0883883461356163 : f32 + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %alpha_raw, %beta_raw, %bias, %a_scale, %state_in, %dst, %state_in, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_rmsnorm_gate_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_rmsnorm_gate_f32_wmma.loom new file mode 100644 index 000000000000..00c2cd64daab --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_scan_rmsnorm_gate_f32_wmma.loom @@ -0,0 +1,41 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +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 @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +kernel.def export("llm_gated_delta_net_f32_wmma_head128_rmsnorm_gate") @llm_gated_delta_net_f32_wmma_head128_rmsnorm_gate() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %dst: buffer, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant false : i1 + %projection_epilogue = scalar.constant false : i1 + %snapshot_stride = index.constant 1 : index + %scale = scalar.constant 0.0883883461356163 : f32 + %rmsnorm_gate = scalar.constant true : i1 + %rms_epsilon = config.get @ggml.rmsnorm_gate_f32.rms_epsilon : f32 + %rms_gate_op = config.get @ggml.rmsnorm_gate_f32.gate_op : index + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %g, %beta, %g, %beta, %state_in, %dst, %state_in, %rmsnorm_gate, %rms_epsilon, %rms_gate_op, %rms_weight, %raw_gate, %norm_output, %half_output, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_selected_snapshot_projection_epilogue_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_selected_snapshot_projection_epilogue_f32_wmma.loom new file mode 100644 index 000000000000..2097217c458b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_selected_snapshot_projection_epilogue_f32_wmma.loom @@ -0,0 +1,39 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.state_row_count : %value: index where [range(%value, 1, 262208)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.snapshot_stride : %value: index where [range(%value, 1, 1073741823)] + +kernel.def export("llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_epilogue") @llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_epilogue() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %state_ids: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant true : i1 + %projection_epilogue = scalar.constant true : i1 + %scale = scalar.constant 0.0883883461356163 : f32 + %snapshot_stride = config.get @llm.gated_delta_net.snapshot_stride : index + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant true : i1 + %state_row_count = config.get @llm.gated_delta_net.state_row_count : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %alpha_raw, %beta_raw, %bias, %a_scale, %state_in, %dst, %snapshot_cache, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %state_ids, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_selected_snapshot_projection_rms_gate_publish_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_selected_snapshot_projection_rms_gate_publish_q8.loom new file mode 100644 index 000000000000..313d86123ffb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_selected_snapshot_projection_rms_gate_publish_q8.loom @@ -0,0 +1,45 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.state_row_count : %value: index where [range(%value, 1, 262208)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.snapshot_stride : %value: index where [range(%value, 1, 1073741823)] + +config.decl @ggml.rmsnorm_gate_f32.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_gate_f32.gate_op : %value: index where [range(%value, 0, 23)] + +kernel.def export("llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_rms_gate_q8") @llm_gated_delta_net_f32_wmma_head128_selected_snapshot_projection_rms_gate_q8() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer, %state_ids: buffer, %rms_weight: buffer, %raw_gate: buffer, %q8_output: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant true : i1 + %projection_epilogue = scalar.constant true : i1 + %scale = scalar.constant 0.0883883461356163 : f32 + %snapshot_stride = config.get @llm.gated_delta_net.snapshot_stride : index + %rms_epsilon = config.get @ggml.rmsnorm_gate_f32.rms_epsilon : f32 + %rms_gate_op = config.get @ggml.rmsnorm_gate_f32.gate_op : index + %rmsnorm_gate = scalar.constant true : i1 + %select_state = scalar.constant true : i1 + %state_row_count = config.get @llm.gated_delta_net.state_row_count : index + %publish_q8 = scalar.constant true : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %alpha_raw, %beta_raw, %bias, %a_scale, %state_in, %dst, %snapshot_cache, %rmsnorm_gate, %rms_epsilon, %rms_gate_op, %rms_weight, %raw_gate, %dst, %q8_output, %select_state, %state_row_count, %state_ids, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_snapshot_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_snapshot_f32_wmma.loom new file mode 100644 index 000000000000..1cfc108a5024 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_snapshot_f32_wmma.loom @@ -0,0 +1,38 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.snapshot_stride : %value: index where [range(%value, 1, 1073741823)] + +// Low-row MTP specialization publishes each recurrence prefix and final state to the reverse-ordered rollback cache. +kernel.def export("llm_gated_delta_net_f32_wmma_head128_snapshot") @llm_gated_delta_net_f32_wmma_head128_snapshot() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant true : i1 + %projection_epilogue = scalar.constant false : i1 + %scale = scalar.constant 0.0883883461356163 : f32 + %snapshot_stride = config.get @llm.gated_delta_net.snapshot_stride : index + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %g, %beta, %g, %beta, %state_in, %dst, %snapshot_cache, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_snapshot_projection_epilogue_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_snapshot_projection_epilogue_f32_wmma.loom new file mode 100644 index 000000000000..b34528b9439e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/head128_snapshot_projection_epilogue_f32_wmma.loom @@ -0,0 +1,39 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.gated_delta_net.f32_wmma_head128.body(%state_inplace: i1, %publish_snapshots: i1, %projection_epilogue: i1, %snapshot_stride: index, %scale: f32, %q: buffer, %k: buffer, %v: buffer, %g: buffer, %beta: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %dst: buffer, %snapshot_cache: buffer, %rmsnorm_gate: i1, %rms_epsilon: f32, %rms_gate_op: index, %rms_weight: buffer, %raw_gate: buffer, %norm_output: buffer, %half_output: buffer, %select_state: i1, %state_row_count: index, %state_ids: buffer, %publish_q8: i1) + +config.decl @llm.gated_delta_net.head_width : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.head_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.sequence_count : %value: index where [range(%value, 1, 4096)] + +config.decl @llm.gated_delta_net.snapshot_stride : %value: index where [range(%value, 1, 1073741823)] + +// Low-row snapshot scan with the alpha/beta projection epilogue folded into +// the existing one-thread-per-token scalar load point. +kernel.def export("llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue") @llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %blk = index.constant 128 : index + %blk_m = index.constant 127 : index + %s_v = config.get @llm.gated_delta_net.head_width : index + %n_heads = config.get @llm.gated_delta_net.head_count : index + %n_seqs = config.get @llm.gated_delta_net.sequence_count : index + %cr = index.add %s_v, %blk_m : index + %col_blocks = index.div %cr, %blk : index + kernel.launch.config workgroups(%n_heads, %n_seqs, %col_blocks) workgroup_size(%wg, %unit, %unit) : index +} launch(%q: buffer, %k: buffer, %v: buffer, %alpha_raw: buffer, %beta_raw: buffer, %bias: buffer, %a_scale: buffer, %state_in: buffer, %snapshot_cache: buffer, %dst: buffer) { + %state_inplace = scalar.constant false : i1 + %publish_snapshots = scalar.constant true : i1 + %projection_epilogue = scalar.constant true : i1 + %scale = scalar.constant 0.0883883461356163 : f32 + %snapshot_stride = config.get @llm.gated_delta_net.snapshot_stride : index + %rmsnorm_gate = scalar.constant false : i1 + %select_state = scalar.constant false : i1 + %state_row_count = index.constant 1 : index + %publish_q8 = scalar.constant false : i1 + template.apply<@llm.gated_delta_net.f32_wmma_head128.body>(%state_inplace, %publish_snapshots, %projection_epilogue, %snapshot_stride, %scale, %q, %k, %v, %alpha_raw, %beta_raw, %bias, %a_scale, %state_in, %dst, %snapshot_cache, %rmsnorm_gate, %scale, %snapshot_stride, %dst, %dst, %dst, %dst, %select_state, %state_row_count, %dst, %publish_q8) : (i1, i1, i1, index, f32, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, i1, f32, index, buffer, buffer, buffer, buffer, i1, index, buffer, i1) + 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 new file mode 100644 index 000000000000..9cc814d4be1f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net/projection_epilogue_f32.loom @@ -0,0 +1,23 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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_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(%element_count: index) { + %unit = index.constant 1 : index + %cneg1 = index.constant -1 : 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(%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_legacy_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom new file mode 100644 index 000000000000..243b31d54dee --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_legacy_f32_f32.loom @@ -0,0 +1,67 @@ +template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_k_matmul_rope_set_rows_legacy_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @llm_attention_k_matmul_rope_set_rows_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %rotated, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @llm_attention_k_matmul_rope_set_rows_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_k_matmul_rope_set_rows_legacy_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_legacy_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @llm.attention_qkv.input_size : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %cache, %indices, %theta, %freq_factors, %cache, %cache) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom new file mode 100644 index 000000000000..b4d982d0d6af --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_tiled_f32_f32.loom @@ -0,0 +1,73 @@ +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) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_k_matmul_rope_set_rows_tiled_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @llm_attention_k_matmul_rope_set_rows_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %rotated, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @llm_attention_k_matmul_rope_set_rows_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_k_matmul_rope_set_rows_tiled_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_tiled_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %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>(%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_k_matmul_rope_set_rows_vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom new file mode 100644 index 000000000000..f606884be090 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_k_matmul_rope_set_rows_vector_f32_f32.loom @@ -0,0 +1,126 @@ +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +amdgpu.target @llm_attention_k_matmul_rope_set_rows_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @llm_attention_k_matmul_rope_set_rows_vector_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %rotated, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @llm_attention_k_matmul_rope_set_rows_vector_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %positions: buffer, %cache: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index + %head_size0 = config.get @llm.attention_qkv.head_size : index + %head_size = index.assume %head_size0 [range(%head_size0, 4, 1024), mul(%head_size0, 4)] : index + %cache_row_count0 = config.get @llm.attention_qkv.cache_row_count : index + %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %publish0 = scalar.andi %publish_pair, %row1_valid : i1 + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %head_channel = index.rem %channel, %head_size : index + %theta_channel = index.div %head_channel, %c2 : index + %half_head_size = index.div %head_size, %c2 : index + %is_f16 = index.cmp eq, %cache_output_format, %c16 : index + %is_f32 = index.cmp eq, %cache_output_format, %c32 : index + %positions_noalias, %indices_noalias, %theta_noalias, %freq_factors_noalias, %cache_noalias = buffer.assume.noalias %positions, %indices, %theta, %freq_factors, %cache : buffer, buffer, buffer, buffer, buffer + %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + %publish1 = scalar.andi %publish0, %valid_index : i1 + %publish = scalar.andi %publish1, %valid_row : i1 + scf.if %publish { + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_scalar = view.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %freq_factors_scalar = view.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %theta_packet = vector.splat %theta_scalar : vector<2xf32> + %freq_factors_packet = vector.splat %freq_factors_scalar : vector<2xf32> + %x_values = vector.from_elements %value0, %c0_f32 : vector<2xf32> + %y_values = vector.from_elements %value1, %c0_f32 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %rx0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 + %ry0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 + %values = vector.from_elements %rx0, %ry0 : vector<2xf32> + %publish_f16 = scalar.andi %publish, %is_f16 : i1 + %publish_f32 = scalar.andi %publish, %is_f32 : i1 + scf.if %publish_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %truncated = vector.fptrunc %values : vector<2xf32> to vector<2xf16> + vector.store %truncated, %cache_view[%cache_row, %channel] : vector<2xf16>, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + vector.store %values, %cache_view[%cache_row, %channel] : vector<2xf32>, view<[%cache_row_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @llm_attention_k_matmul_rope_set_rows_vector_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_k_matmul_rope_set_rows_vector_gfx11_wave64) @llm_attention_k_matmul_rope_set_rows_vector_f32_f32(%token_count: index) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %indices: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %positions, %cache, %indices, %theta, %freq_factors, %cache, %cache, %cache) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom new file mode 100644 index 000000000000..2de72456e220 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_legacy_f32_f32.loom @@ -0,0 +1,61 @@ +template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_query_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @llm_attention_q_matmul_rope_legacy_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @llm_attention_q_matmul_rope_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_query_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %rotated, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @llm_attention_q_matmul_rope_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_q_matmul_rope_legacy_gfx11_wave64) @llm_attention_q_matmul_rope_legacy_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @llm.attention_qkv.input_size : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_tile, %token_tile, %input, %weight, %positions, %output, %theta, %freq_factors, %output, %output, %output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom new file mode 100644 index 000000000000..e70144aa4d02 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_tiled_f32_f32.loom @@ -0,0 +1,67 @@ +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) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_query_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @llm_attention_q_matmul_rope_tiled_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @llm_attention_q_matmul_rope_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_query_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %rotated, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @llm_attention_q_matmul_rope_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_q_matmul_rope_tiled_gfx11_wave64) @llm_attention_q_matmul_rope_tiled_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %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>(%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_q_matmul_rope_vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom new file mode 100644 index 000000000000..d936e62a7e13 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_q_matmul_rope_vector_f32_f32.loom @@ -0,0 +1,104 @@ +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @llm.attention_qkv.rope_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: vector<4xf32>) -> (vector<4xf32>) + +template.decl @llm.attention_qkv.store_query_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +func.decl @ggml_rope_f32_pair_packet(%position: f32, %theta: vector<2xf32>, %freq_factors: vector<2xf32>, %x_values: vector<2xf32>, %y_values: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + +amdgpu.target @llm_attention_q_matmul_rope_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] + +config.decl @llm.attention_qkv.head_count : %value: index where [range(%value, 1, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @llm_attention_q_matmul_rope_vector_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %head_size = config.get @llm.attention_qkv.head_size : index + %head_count = config.get @llm.attention_qkv.head_count : index + scf.if %publish_word { + %rotated = template.apply<@llm.attention_qkv.rope_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %positions, %theta, %freq_factors, %values) : (index, index, index, index, index, buffer, buffer, buffer, vector<4xf32>) -> (vector<4xf32>) + template.apply<@llm.attention_qkv.store_query_vector4>(%token0, %channel, %token_count0, %head_size, %head_count, %rotated, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @llm_attention_q_matmul_rope_vector_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %positions: buffer, %output: buffer, %theta: buffer, %freq_factors: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 4, 32768), mul(%output_size0, 4)] : index + %head_size0 = config.get @llm.attention_qkv.head_size : index + %head_count0 = config.get @llm.attention_qkv.head_count : index + %head_size, %head_count = index.assume %head_size0, %head_count0 [range(%head_size0, 4, 1024), mul(%head_size0, 4), range(%head_count0, 1, 64)] : index, index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %publish = scalar.andi %publish_pair, %row1_valid : i1 + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %head = index.div %channel, %head_size : index + %head_channel = index.rem %channel, %head_size : index + %theta_channel = index.div %head_channel, %c2 : index + %half_head_size = index.div %head_size, %c2 : index + %positions_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer + scf.if %publish { + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%head_count]x[%head_size]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_scalar = view.load %theta_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %freq_factors_scalar = view.load %freq_factors_view[%theta_channel] : view<[%half_head_size]xf32> -> f32 + %theta_packet = vector.splat %theta_scalar : vector<2xf32> + %freq_factors_packet = vector.splat %freq_factors_scalar : vector<2xf32> + %x_values = vector.from_elements %value0, %c0_f32 : vector<2xf32> + %y_values = vector.from_elements %value1, %c0_f32 : vector<2xf32> + %rotated_x, %rotated_y = func.call @ggml_rope_f32_pair_packet(%position, %theta_packet, %freq_factors_packet, %x_values, %y_values) : (f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %rx0 = vector.extract %rotated_x[0] : vector<2xf32> -> f32 + %ry0 = vector.extract %rotated_y[0] : vector<2xf32> -> f32 + %values = vector.from_elements %rx0, %ry0 : vector<2xf32> + vector.store %values, %output_view[%token, %head, %head_channel] : vector<2xf32>, view<[%launch_token_count]x[%head_count]x[%head_size]xf32> + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @llm_attention_q_matmul_rope_vector_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_q_matmul_rope_vector_gfx11_wave64) @llm_attention_q_matmul_rope_vector_f32_f32(%token_count: index) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %positions: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %positions, %output, %theta, %freq_factors, %output, %output, %output, %output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom new file mode 100644 index 000000000000..ff5015d47fb0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_legacy_f32_f32.loom @@ -0,0 +1,58 @@ +template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_v_matmul_set_rows_legacy_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @llm_attention_v_matmul_set_rows_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %values, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @llm_attention_v_matmul_set_rows_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_v_matmul_set_rows_legacy_gfx11_wave64) @llm_attention_v_matmul_set_rows_legacy_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %input_size = config.get @llm.attention_qkv.input_size : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_tile, %token_tile, %input, %weight, %indices, %cache, %cache, %cache, %cache, %cache, %cache) : (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 new file mode 100644 index 000000000000..129beb52a9f3 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_tiled_f32_f32.loom @@ -0,0 +1,66 @@ +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) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_v_matmul_set_rows_tiled_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat.activation_format : %value: index where [range(%value, 0, 16)] + +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] + +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4> device @llm_attention_v_matmul_set_rows_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %values, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @llm_attention_v_matmul_set_rows_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_v_matmul_set_rows_tiled_gfx11_wave64) @llm_attention_v_matmul_set_rows_tiled_f32_f32(%token_count: index) { + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %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>(%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/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom new file mode 100644 index 000000000000..6c8bded738f4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/llm/llm_attention_v_matmul_set_rows_vector_f32_f32.loom @@ -0,0 +1,108 @@ +template.decl @ggml.matmul.vector_activation.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @llm.attention_qkv.store_cache_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: buffer, %arg7: vector<4xf32>, %arg8: buffer) + +amdgpu.target @llm_attention_v_matmul_set_rows_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @llm.attention_qkv.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] +config.decl @llm.attention_qkv.output_size : %value: index where [range(%value, 1, 32768)] +config.decl @llm.attention_qkv.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @llm.attention_qkv.cache_row_count : %value: index where [range(%value, 1, 1048576)] +config.decl @llm.attention_qkv.cache_output_format : %value: index where [range(%value, 16, 32)] +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @llm_attention_v_matmul_set_rows_vector_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %cache_row_count = config.get @llm.attention_qkv.cache_row_count : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + scf.if %publish_word { + template.apply<@llm.attention_qkv.store_cache_vector4>(%cache_output_format, %token0, %channel, %token_count0, %cache_row_count, %output_size0, %indices, %values, %cache) : (index, index, index, index, index, index, buffer, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @llm_attention_v_matmul_set_rows_vector_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %indices: buffer, %cache: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 1)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 32768)] : index + %cache_row_count0 = config.get @llm.attention_qkv.cache_row_count : index + %cache_row_count = index.assume %cache_row_count0 [range(%cache_row_count0, 1, 1048576)] : index + %cache_output_format = config.get @llm.attention_qkv.cache_output_format : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_i64 = scalar.constant 0 : i64 + %c1048575_i64 = scalar.constant 1048575 : i64 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %is_f16 = index.cmp eq, %cache_output_format, %c16 : index + %is_f32 = index.cmp eq, %cache_output_format, %c32 : index + %indices_noalias, %cache_noalias = buffer.assume.noalias %indices, %cache : buffer, buffer + %indices_view = buffer.view %indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> + %index_raw = view.load %indices_view[%token] : view<[%launch_token_count]xi64> -> i64 + %index_nonnegative = scalar.cmpi sge, %index_raw, %c0_i64 : i64 + %index_in_cast_range = scalar.cmpi sle, %index_raw, %c1048575_i64 : i64 + %valid_index = scalar.andi %index_nonnegative, %index_in_cast_range : i1 + %safe_index0_i64 = scf.select %valid_index, %index_raw, %c0_i64 : i64 + %safe_index_i64 = scalar.assume %safe_index0_i64 [range(%safe_index0_i64, 0, 1048575)] : i64 + %cache_row0 = index.cast %safe_index_i64 : i64 to index + %valid_row = index.cmp ult, %cache_row0, %cache_row_count : index + %cache_row = scf.select %valid_row, %cache_row0, %c0 : index + %publish0 = scalar.andi %publish_pair, %valid_index : i1 + %publish1 = scalar.andi %publish0, %valid_row : i1 + %publish_first_f16 = scalar.andi %publish1, %is_f16 : i1 + %publish_first_f32 = scalar.andi %publish1, %is_f32 : i1 + %publish_second = scalar.andi %publish1, %row1_valid : i1 + %publish_second_f16 = scalar.andi %publish_second, %is_f16 : i1 + %publish_second_f32 = scalar.andi %publish_second, %is_f32 : i1 + scf.if %publish_first_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %truncated = scalar.fptrunc %value0 : f32 to f16 + view.store %truncated, %cache_view[%cache_row, %channel] : f16, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_first_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + view.store %value0, %cache_view[%cache_row, %channel] : f32, view<[%cache_row_count]x[%output_size]xf32> + } + scf.if %publish_second_f16 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf16> + %channel1 = index.add %channel, %c1 : index + %truncated = scalar.fptrunc %value1 : f32 to f16 + view.store %truncated, %cache_view[%cache_row, %channel1] : f16, view<[%cache_row_count]x[%output_size]xf16> + } + scf.if %publish_second_f32 { + %cache_view = buffer.view %cache_noalias[%c0_offset] : buffer -> view<[%cache_row_count]x[%output_size]xf32> + %channel1 = index.add %channel, %c1 : index + view.store %value1, %cache_view[%cache_row, %channel1] : f32, view<[%cache_row_count]x[%output_size]xf32> + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @llm_attention_v_matmul_set_rows_vector_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@llm_attention_v_matmul_set_rows_vector_gfx11_wave64) @llm_attention_v_matmul_set_rows_vector_f32_f32(%token_count: index) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %indices: buffer, %cache: buffer) { + %input_size = config.get @llm.attention_qkv.input_size : index + %output_size = config.get @llm.attention_qkv.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @llm.attention_qkv.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %indices, %cache, %cache, %cache, %cache, %cache, %cache, %cache) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/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/skinny_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom new file mode 100644 index 000000000000..7f5839752649 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/skinny_input_f32_publish_f32.loom @@ -0,0 +1,144 @@ +// Skinny GGML matmul for contiguous 2D F32 activations. +// +// The accumulation core covers low-token direct F32 and staged quantized +// activation paths. This wrapper selects the output epilogue. +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.skinny_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.skinny_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.skinny_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.mul_mat.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_postops.epilogue : %value: index where [range(%value, 1, 3)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 2, 5)] + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_skinny_input_f32_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 2, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %uses_bias_only = index.cmp eq, %epilogue, %c1 : index + %uses_residual_only = index.cmp eq, %epilogue, %c2 : index + %uses_bias_residual = index.cmp eq, %epilogue, %c3 : index + %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 + %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %accumulated = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.skinny_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %accumulated, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %accumulated : vector<4xf32> + } + %residual_added = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.skinny_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + %published = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %residual_added) : (index, vector<4xf32>) -> (vector<4xf32>) + scf.if %uses_residual { + template.apply<@ggml.matmul.skinny_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.skinny_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + template.return +} + +template.def<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.finish_tile> device @ggml_matmul_skinny_input_f32_publish_f32_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat.input_size : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 5), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + + + +kernel.def target(@ggml_matmul_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = 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, 2, 5), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.matmul.skinny_input_f32_dequant_accumulate_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %identity, %epilogue, %c0_f32, %channel_tile, %token_tile, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/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 new file mode 100644 index 000000000000..71fca22091be --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_input_f32_publish_f32.loom @@ -0,0 +1,289 @@ +// Tiled GGML matmul for contiguous 2D F32 activations. +// +// This wrapper owns the single-output tiled kernel family. The tiled +// accumulation core is shared; each kernel symbol selects an epilogue variant. +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.core(%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) + +template.decl @ggml.matmul.tiled_input_f32_dequant_accumulate_f32.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %epilogue: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_publish.store_vector4(%format: index, %publish_word: i1, %token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %scratch_values: view<256xf32>, %scratch_d: view<32xf32>, %output: buffer) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.mul_mat.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat.activation_format : %value: index where [range(%value, 0, 16)] + +config.decl @ggml.mul_mat_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 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) { + %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 + %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> + %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 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %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_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 { + scf.yield %values : vector<4xf32> + } + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %accumulated, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %accumulated : vector<4xf32> + } + %residual_added = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + %published = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %residual_added) : (index, vector<4xf32>) -> (vector<4xf32>) + scf.if %publishes_f16_alternate { + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> + %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%c32, %publish_word, %token_count0, %output_size0, %token0, %channel, %published, %scratch_values, %scratch_d, %output) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + template.apply<@ggml.matmul.tiled_publish.store_vector4>(%c16, %publish_word, %token_count0, %output_size0, %token0, %channel, %published, %scratch_values, %scratch_d, %residual_input) : (index, i1, index, index, index, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } else { + scf.if %uses_residual { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + } + template.return +} + +template.def<@ggml.matmul.tiled_input_f32_dequant_accumulate_f32.finish_tile> device @ggml_matmul_tiled_input_f32_publish_f32_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 0 : index + %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>(%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 +} + +kernel.def target(@ggml_matmul_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %f16_output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat.output_accumulation : index + %output_unary_op = config.get @ggml.mul_mat.output_unary_op : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %epilogue = index.constant 4 : index + %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>(%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 + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) { + %input_size = config.get @ggml.mul_mat_postops.input_size : index + %output_size = config.get @ggml.mul_mat_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.mul_mat_postops.weight_format : index + %c0 = index.constant 0 : index + %identity = index.constant 23 : index + %epilogue = 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>(%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 new file mode 100644 index 000000000000..d50185e7dbbd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32.loom @@ -0,0 +1,176 @@ +// Generic paired tiled GGML matmul for contiguous 2D F32 activations. +// This wrapper composes paired tiled accumulation with pair postprocess, +// optional epilogue adds, and F32 publication. +template.decl @ggml.matmul.tiled_pair_input_f32_dequant_accumulate_f32.core(%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) + +template.decl @ggml.matmul.tiled_pair_postprocess_f32.apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_bias_vector4(%output_size0: index, %channel: index, %values: vector<4xf32>, %bias: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_postprocess_f32.add_residual_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %residual_input: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +amdgpu.target @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_swiglu.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_swiglu.gate_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_swiglu.up_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_swiglu.op : %value: index where [range(%value, 0, 8)] + +config.decl @ggml.mul_mat_swiglu.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_pair_input_f32_dequant_accumulate_f32.publish_vector4> device @ggml_matmul_tiled_pair_input_f32_binary_publish_f32_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %op: index, %epilogue: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %uses_bias_only = index.cmp eq, %epilogue, %c1 : index + %uses_residual_only = index.cmp eq, %epilogue, %c2 : index + %uses_bias_residual = index.cmp eq, %epilogue, %c3 : index + %uses_bias = scalar.ori %uses_bias_only, %uses_bias_residual : i1 + %uses_residual = scalar.ori %uses_residual_only, %uses_bias_residual : i1 + scf.if %publish_word { + %binary = template.apply<@ggml.matmul.tiled_pair_postprocess_f32.apply_vector4>(%op, %lhs, %rhs) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %biased = scf.if %uses_bias -> (vector<4xf32>) { + %with_bias = template.apply<@ggml.matmul.tiled_postprocess_f32.add_bias_vector4>(%output_size0, %channel, %binary, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_bias : vector<4xf32> + } else { + scf.yield %binary : vector<4xf32> + } + %published = scf.if %uses_residual -> (vector<4xf32>) { + %with_residual = template.apply<@ggml.matmul.tiled_postprocess_f32.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %biased, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %with_residual : vector<4xf32> + } else { + scf.yield %biased : vector<4xf32> + } + scf.if %uses_residual { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + } else { + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + } + } + template.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_gfx11_wave64) @ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 0 : index + %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>(%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 + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = 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>(%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 new file mode 100644 index 000000000000..78f81d9d6992 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/tiled_pair_input_f32_binary_publish_f32_k16.loom @@ -0,0 +1,65 @@ +// Generic paired tiled GGML matmul with F32 and K16-major F16 publication. +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) + +template.decl @ggml.matmul.tiled_pair_postprocess_f32.apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.matmul.tiled_publish_f32.store_vector4(%token_count0: index, %output_size0: index, %token0: index, %channel: index, %values: vector<4xf32>, %output: buffer) + +template.decl @ggml.normalization.publish_f16.k16_vector4(%values: vector<4xf32>, %channel: index, %token: index, %token_count: index, %hidden_size: index, %packet: index, %output: buffer) + +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) { + %c4 = index.constant 4 : index + scf.if %publish_word { + %published = template.apply<@ggml.matmul.tiled_pair_postprocess_f32.apply_vector4>(%op, %lhs, %rhs) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + template.apply<@ggml.matmul.tiled_publish_f32.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %output) : (index, index, index, index, vector<4xf32>, buffer) + %packet = index.div %channel, %c4 : index + template.apply<@ggml.normalization.publish_f16.k16_vector4>(%published, %channel, %token0, %token_count0, %output_size0, %packet, %residual_output) : (vector<4xf32>, index, index, index, index, index, buffer) + } + template.return +} + +kernel.def target(@ggml_matmul_tiled_pair_input_f32_binary_publish_f32_k16_gfx11_wave64) export("ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16") @ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %epilogue = index.constant 0 : index + %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>(%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_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_f32_f32.loom new file mode 100644 index 000000000000..0f87ac2802b4 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_f32_f32.loom @@ -0,0 +1,126 @@ +// Vector-shaped GGML matmul reduction for batch-size-1 F32 outputs. +template.decl @ggml.matmul.vector_publish.next_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) +template.decl @ggml.matmul.vector_activation.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.unary_f32.apply(%arg0: index, %arg1: f32) -> (f32) + +template.decl @ggml.unary_f32.apply_vector4(%arg0: index, %arg1: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_matmul_vector_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 1, 524288)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector.output_unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @ggml.matmul.vector.weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @ggml_matmul_vector_f32_f32_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %output_noalias = buffer.assume.noalias %output : buffer + %partition = kernel.workgroup.id : index + %element_bytes = index.constant 4 : offset + %plane_elements = index.mul %token_count, %output_size : index + %plane_bytes = index.scale %plane_elements, %element_bytes : index, offset -> offset + %plane_offset = index.scale %partition, %plane_bytes : index, offset -> offset + %output_view = buffer.view %output_noalias[%plane_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_word { + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + %published = scf.if %accumulates_output -> (vector<4xf32>) { + %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%output_size]xf32>, vector<4xi1>, vector<4xf32> + %sum = vector.addf %residual, %values : vector<4xf32> + scf.yield %sum : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %activated = template.apply<@ggml.unary_f32.apply_vector4>(%output_unary_op, %published) : (index, vector<4xf32>) -> (vector<4xf32>) + vector.store.mask %activated, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%output_size]xf32>, vector<4xi1> + template.apply<@ggml.matmul.vector_publish.next_vector4>(%publish_word, %publish_format, %token_count0, %output_size0, %token0, %channel, %activated, %next_output) : (i1, index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @ggml_matmul_vector_f32_f32_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + scf.if %publish_pair { + %published0 = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %residual, %value0 : f32 + scf.yield %sum : f32 + } else { + scf.yield %value0 : f32 + } + %activated0 = template.apply<@ggml.unary_f32.apply>(%output_unary_op, %published0) : (index, f32) -> (f32) + view.store %activated0, %output_view[%token, %channel] : f32, view<[%token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %published1 = scf.if %accumulates_output -> (f32) { + %residual = view.load %output_view[%token, %row1] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %residual, %value1 : f32 + scf.yield %sum : f32 + } else { + scf.yield %value1 : f32 + } + %activated1 = template.apply<@ggml.unary_f32.apply>(%output_unary_op, %published1) : (index, f32) -> (f32) + view.store %activated1, %output_view[%token, %row1] : f32, view<[%token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @ggml_matmul_vector_f32_f32_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_matmul_vector_gfx11_wave64) @ggml_mul_mat_vector_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_accumulation = config.get @ggml.matmul.vector.output_accumulation : index + %output_unary_op = config.get @ggml.matmul.vector.output_unary_op : index + %weight_format = config.get @ggml.matmul.vector.weight_format : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_group, %token_group, %input, %weight, %output, %output, %output, %output, %output, %output, %output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom new file mode 100644 index 000000000000..65402e073a27 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_postops_f32_f32.loom @@ -0,0 +1,135 @@ +template.decl @ggml.matmul.vector_activation.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul.vector_activation.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) -> (vector<4xf32>) +template.decl @ggml.matmul.vector_activation.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: vector<4xf32>, %arg5: buffer) + +template.decl @ggml.matmul.vector_publish.next_vector4(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) +template.decl @ggml.matmul.vector_activation.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer) + +template.decl @ggml.matmul.vector_activation.finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul.vector_activation.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.matmul.vector_activation.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +template.decl @ggml.matmul.vector_activation.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) + +amdgpu.target @ggml_matmul_vector_postops_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.matmul.vector_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.matmul.vector_postops.output_size : %value: index where [range(%value, 1, 524288)] + +config.decl @ggml.matmul.vector_postops.weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.matmul.vector_postops.apply_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.matmul.vector_postops.apply_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +template.def<@ggml.matmul.vector_activation.publish_vector4> device @ggml_mul_mat_vector_postops_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %apply_bias0 = config.get @ggml.matmul.vector_postops.apply_bias : index + %apply_residual0 = config.get @ggml.matmul.vector_postops.apply_residual : index + %c1 = index.constant 1 : index + %apply_bias = index.cmp eq, %apply_bias0, %c1 : index + %apply_residual = index.cmp eq, %apply_residual0, %c1 : index + scf.if %publish_word { + %with_bias = scf.if %apply_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul.vector_activation.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %apply_residual -> (vector<4xf32>) { + %residual = template.apply<@ggml.matmul.vector_activation.add_residual_vector4>(%token_count0, %output_size0, %token0, %channel, %with_bias, %residual_input) : (index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %residual : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul.vector_activation.store_vector4>(%token_count0, %output_size0, %token0, %channel, %published, %residual_output) : (index, index, index, index, vector<4xf32>, buffer) + template.apply<@ggml.matmul.vector_publish.next_vector4>(%publish_word, %publish_format, %token_count0, %output_size0, %token0, %channel, %published, %next_output) : (i1, index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul.vector_activation.publish_pair2> device @ggml_mul_mat_vector_postops_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output_accumulation: index, %output_unary_op: index, %publish_format: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer, %next_output: buffer) { + %apply_bias0 = config.get @ggml.matmul.vector_postops.apply_bias : index + %apply_residual0 = config.get @ggml.matmul.vector_postops.apply_residual : index + %token_count = index.assume %token_count0 [range(%token_count0, 1, 5)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 524288)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %apply_bias = index.cmp eq, %apply_bias0, %c1 : index + %apply_residual = index.cmp eq, %apply_residual0, %c1 : index + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %bias_noalias, %residual_input_noalias, %residual_output_noalias = buffer.assume.noalias %bias, %residual_input, %residual_output : buffer, buffer, buffer + %bias_view = buffer.view %bias_noalias[%c0_offset] : buffer -> view<[%output_size]xf32> + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + %residual_output_view = buffer.view %residual_output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%output_size]xf32> + scf.if %publish_pair { + %biased0 = scf.if %apply_bias -> (f32) { + %bias_value = view.load %bias_view[%channel] : view<[%output_size]xf32> -> f32 + %sum = scalar.addf %value0, %bias_value : f32 + scf.yield %sum : f32 + } else { + scf.yield %value0 : f32 + } + %published0 = scf.if %apply_residual -> (f32) { + %residual = view.load %residual_input_view[%token, %channel] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %biased0, %residual : f32 + scf.yield %sum : f32 + } else { + scf.yield %biased0 : f32 + } + view.store %published0, %residual_output_view[%token, %channel] : f32, view<[%token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %biased1 = scf.if %apply_bias -> (f32) { + %bias_value = view.load %bias_view[%row1] : view<[%output_size]xf32> -> f32 + %sum = scalar.addf %value1, %bias_value : f32 + scf.yield %sum : f32 + } else { + scf.yield %value1 : f32 + } + %published1 = scf.if %apply_residual -> (f32) { + %residual = view.load %residual_input_view[%token, %row1] : view<[%token_count]x[%output_size]xf32> -> f32 + %sum = scalar.addf %biased1, %residual : f32 + scf.yield %sum : f32 + } else { + scf.yield %biased1 : f32 + } + view.store %published1, %residual_output_view[%token, %row1] : f32, view<[%token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.matmul.vector_activation.finish_output_group> device @ggml_mul_mat_vector_postops_finish_output_group(%token_count: index, %output_size: index, %output_group_count: index, %token_group_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_group: index, %token_group: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + + + +kernel.def target(@ggml_matmul_vector_postops_gfx11_wave64) @ggml_mul_mat_vector_bias_residual_f32_f32(%token_count: index) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c1 = index.constant 1 : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.matmul.vector_activation.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer, %next_output: buffer) { + %input_size = config.get @ggml.matmul.vector_postops.input_size : index + %output_size = config.get @ggml.matmul.vector_postops.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %weight_format = config.get @ggml.matmul.vector_postops.weight_format : index + %c0 = index.constant 0 : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1), le(%token_count, %token_capacity)] : index + %channel_group = kernel.workgroup.id : index + %token_group = kernel.workgroup.id : index + template.apply<@ggml.matmul.vector_activation.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %c0, %c0, %c0_f32, %channel_group, %token_group, %input, %weight, %bias, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output, %residual_output, %next_output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_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/matmul/vector_q6_f32_f32_kernel.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_f32_f32_kernel.loom new file mode 100644 index 000000000000..870419b2b38b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_f32_f32_kernel.loom @@ -0,0 +1,39 @@ +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.def @ggml.matmul.vector.packed_layout = 1 : index + +config.decl @ggml.matmul.vector.q6_storage : %value: index where [range(%value, 0, 1)] + +template.decl @ggml.matmul.vector_q6_f32.body_packed_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_q6_f32.body_i8_storage(%token_count: index, %input_size0: index, %output_size0: index, %output_block: index, %input: buffer, %weight: buffer, %output: buffer) + +kernel.def export("ggml_mul_mat_vector_q6_f32_f32") @ggml_mul_mat_vector_q6_f32_f32(%token_count: index) { + %output_size = config.get @ggml.matmul.vector.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_blocks = index.div %padded_output_size, %c64 : index + kernel.launch.config workgroups(%output_blocks, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size = config.get @ggml.matmul.vector.output_size : index + %scale_row_layout = config.get @ggml.matmul.vector.packed_layout : index + %q6_storage = config.get @ggml.matmul.vector.q6_storage : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %output_block = kernel.workgroup.id : index + %c1 = index.constant 1 : index + %uses_i8_storage = index.cmp eq, %q6_storage, %c1 : index + scf.if %uses_i8_storage { + template.apply<@ggml.matmul.vector_q6_f32.body_i8_storage>(%bounded_token_count, %input_size, %output_size, %output_block, %input, %weight, %output) : (index, index, index, index, buffer, buffer, buffer) + } else { + template.apply<@ggml.matmul.vector_q6_f32.body_packed_storage>(%bounded_token_count, %input_size, %output_size, %output_block, %scale_row_layout, %input, %weight, %output) : (index, index, index, index, index, buffer, buffer, buffer) + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom new file mode 100644 index 000000000000..d8ca1ee69144 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul/vector_q6_selected_outputs_f32_f32_kernel.loom @@ -0,0 +1,30 @@ +config.decl @ggml.matmul.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.matmul.vector.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.matmul.vector.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.def @ggml.matmul.vector.packed_layout = 1 : index + +config.decl @ggml.matmul.vector.weight_offset : %value: index where [range(%value, 0, 2147483647), mul(%value, 256)] + +template.decl @ggml.matmul.vector_q6_f32.body_selected_outputs(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %selected_count0: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %output: buffer) + +template.decl @ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32(%selected_count: index, %output_size: index, %selected_outputs: buffer, %values: buffer, %output: buffer) + +kernel.def export("ggml_mul_mat_vector_q6_selected_outputs_f32_f32") @ggml_mul_mat_vector_q6_selected_outputs_f32_f32(%token_count: index, %selected_count: index) { + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %selected_count: index, %input: buffer, %weight: buffer, %selected_outputs: buffer, %selected_values: buffer, %output: buffer) { + %input_size = config.get @ggml.matmul.vector.input_size : index + %output_size0 = config.get @ggml.matmul.vector.output_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %scale_row_layout = config.get @ggml.matmul.vector.packed_layout : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_selected_count = index.assume %selected_count [range(%selected_count, 1, 64)] : index + template.apply<@ggml.matmul.vector_q6_f32.body_selected_outputs>(%bounded_token_count, %input_size, %output_size, %scale_row_layout, %bounded_selected_count, %input, %weight, %selected_outputs, %selected_values) : (index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + template.apply<@ggml.matmul.vector_publish.f32_scatter_selected_outputs_wave32>(%bounded_selected_count, %output_size, %selected_outputs, %selected_values, %output) : (index, index, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mul_mat_id_f16_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mul_mat_id_f16_f16_wmma.loom new file mode 100644 index 000000000000..af160bd36eaf --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mul_mat_id_f16_f16_wmma.loom @@ -0,0 +1,50 @@ +// Standalone F16 MUL_MAT_ID fallback for routed expert projections. + +template.decl @ggml.mul_mat_id_f16_f16_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer) + +template.decl @ggml.mul_mat_id_f16_f16_wmma.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %channel_tile: index, %expert: index, %route_tile_base: index, %expert_route_count: index, %expert_table: buffer, %output: buffer) + +template.decl @ggml.mul_mat_id_f16_f16_wmma.publish_vector4(%publish_word: i1, %assignment: index, %channel: index, %token_count: index, %route_count: index, %output_size: index, %values: vector<4xf16>, %output: buffer) + +amdgpu.target @ggml_mul_mat_id_f16_f16_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_f16_f16.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_f16_f16.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @ggml.mul_mat_id_f16_f16.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_f16_f16.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @ggml.mul_mat_id_f16_f16.weight_format : %value: index where [range(%value, 4, 6)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +// Q4_K and Q6_K share one entry point. The model dispatcher supplies a +// compile-time weight format so specialization erases the inactive decoder. +kernel.def target(@ggml_mul_mat_id_f16_f16_gfx11_wave64) @ggml_mul_mat_id_f16_f16_wmma(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_f16_f16.expert_count : index + %output_size = config.get @ggml.mul_mat_id_f16_f16.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_count, %c63 : index + %route_tiles = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_f16_f16.input_size : index + %route_count = config.get @ggml.mul_mat_id_f16_f16.route_count : index + %expert_count = config.get @ggml.mul_mat_id_f16_f16.expert_count : index + %output_size = config.get @ggml.mul_mat_id_f16_f16.output_size : index + %weight_format = config.get @ggml.mul_mat_id_f16_f16.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %route_tile = kernel.workgroup.id : index + %expert = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_id_f16_f16_wmma.core>(%weight_format, %bounded_token_count, %input_size, %route_count, %expert_count, %output_size, %channel_tile, %route_tile, %expert, %input, %expert_table, %weight, %output) : (index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mul_mat_id_swiglu_f16_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mul_mat_id_swiglu_f16_f16_wmma.loom new file mode 100644 index 000000000000..1643dbf64615 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/mul_mat_id_swiglu_f16_f16_wmma.loom @@ -0,0 +1,57 @@ +// Routed gate/up SwiGLU F16-output path with compile-time weight-format selection. + +template.decl @ggml.mul_mat_id_swiglu_f16_f16.body(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: i32, %arg8: i32, %arg9: i32, %arg10: index, %arg11: index, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) + +amdgpu.target @ggml_mul_mat_id_swiglu_f16_f16_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.gate_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.up_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.descriptor_expert_mask : %value: i32 + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.descriptor_partition_shift : %value: i32 + +config.decl @ggml.mul_mat_id_swiglu_f16_f16.descriptor_row_count_shift : %value: i32 + +kernel.def target(@ggml_mul_mat_id_swiglu_f16_f16_gfx11_wave32) @ggml_mul_mat_id_swiglu_f16_f16_wmma(%token_count: index) { + %route_count = config.get @ggml.mul_mat_id_swiglu_f16_f16.route_count : index + %expert_count = config.get @ggml.mul_mat_id_swiglu_f16_f16.expert_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu_f16_f16.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_swiglu_f16_f16.input_size : index + %route_count = config.get @ggml.mul_mat_id_swiglu_f16_f16.route_count : index + %expert_count = config.get @ggml.mul_mat_id_swiglu_f16_f16.expert_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu_f16_f16.output_size : index + %gate_weight_format = config.get @ggml.mul_mat_id_swiglu_f16_f16.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_id_swiglu_f16_f16.up_weight_format : index + %descriptor_expert_mask = config.get @ggml.mul_mat_id_swiglu_f16_f16.descriptor_expert_mask : i32 + %descriptor_partition_shift = config.get @ggml.mul_mat_id_swiglu_f16_f16.descriptor_partition_shift : i32 + %descriptor_row_count_shift = config.get @ggml.mul_mat_id_swiglu_f16_f16.descriptor_row_count_shift : i32 + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_id_swiglu_f16_f16.body>(%gate_weight_format, %up_weight_format, %token_count, %input_size, %route_count, %expert_count, %output_size, %descriptor_expert_mask, %descriptor_partition_shift, %descriptor_row_count_shift, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, i32, i32, i32, index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom new file mode 100644 index 000000000000..7a061772f1f5 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_next_rmsnorm_publish_f32.loom @@ -0,0 +1,98 @@ +template.decl @ggml.matmul_id.skinny_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.mul_mat_id_postops.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_postops.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_postops.has_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.rms_epsilon : f32 + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.skinny_f32_f32.publish_vector4> device @ggml_mul_mat_id_skinny_postops_next_rmsnorm_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index + %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index + %c0 = index.constant 0 : index + %has_bias = index.cmp ne, %has_bias_value, %c0 : index + %has_residual = index.cmp ne, %has_residual_value, %c0 : index + scf.if %publish_word { + %with_bias = scf.if %has_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul_id.skinny_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %has_residual -> (vector<4xf32>) { + %added = template.apply<@ggml.matmul_id.skinny_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %added : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul_id.skinny_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %residual_output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_tile> device @ggml_mul_mat_id_skinny_postops_next_rmsnorm_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.apply<@ggml.matmul_id.skinny_f32_f32.finish_routed_rmsnorm_weight>(%token_count, %route_count, %output_size, %expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %descriptor_ordinal, %expert, %route_tile_base, %partition_row_count, %expert_table, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_postops.input_size : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_postops.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_id_postops.weight_format : index + %epsilon = config.get @ggml.mul_mat_id_postops.rms_epsilon : f32 + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %epsilon, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom new file mode 100644 index 000000000000..d35eaee7b34e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_postops_publish_f32.loom @@ -0,0 +1,93 @@ +template.decl @ggml.matmul_id.skinny_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.skinny_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_input_f32_postops_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id_postops.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_postops.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_postops.has_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.skinny_f32_f32.publish_vector4> device @ggml_mul_mat_id_skinny_postops_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index + %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index + %c0 = index.constant 0 : index + %has_bias = index.cmp ne, %has_bias_value, %c0 : index + %has_residual = index.cmp ne, %has_residual_value, %c0 : index + scf.if %publish_word { + %with_bias = scf.if %has_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul_id.skinny_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %has_residual -> (vector<4xf32>) { + %added = template.apply<@ggml.matmul_id.skinny_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %added : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul_id.skinny_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_tile> device @ggml_mul_mat_id_skinny_postops_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_input_f32_postops_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_input_f32_postops_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_postops.input_size : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_postops.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_id_postops.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %output, %residual_input, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom new file mode 100644 index 000000000000..a0f41b4dd7eb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_input_f32_publish_f32.loom @@ -0,0 +1,81 @@ +// Standalone GGML MUL_MAT_ID fallback for routed expert projections. +// +// The kernel consumes a precomputed expert table. Each table row stores compact +// assignment indices encoded as token * route_count + route. Input can either +// be broadcast across all routes or contain one input row per route. +template.decl @ggml.matmul_id.skinny_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.skinny_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.skinny_f32_f32.publish_vector4> device @ggml_mul_mat_id_skinny_f32_f32_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %output_row_count = index.mul %token_count, %route_count : index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%output_size]xf32> + scf.if %publish_word { + %bounded_assignment, %assignment_count = index.assume %assignment, %output_row_count [lt(%assignment, %output_row_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<4xf32>, view<[%output_row_count]x[%output_size]xf32>, vector<4xi1> + } + template.return +} + +template.def<@ggml.matmul_id.skinny_f32_f32.finish_tile> device @ggml_mul_mat_id_skinny_f32_f32_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_mul_mat_id_skinny_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_input_f32_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id.expert_count : index + %route_count = config.get @ggml.mul_mat_id.route_count : index + %output_size = config.get @ggml.mul_mat_id.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id.input_size : index + %route_count = config.get @ggml.mul_mat_id.route_count : index + %input_route_count = config.get @ggml.mul_mat_id.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id.expert_count : index + %output_size = config.get @ggml.mul_mat_id.output_size : index + %weight_format = config.get @ggml.mul_mat_id.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.skinny_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom new file mode 100644 index 000000000000..8fd00c85207a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/skinny_pair_input_f32_swiglu_publish_f32.loom @@ -0,0 +1,54 @@ +// Routed GGML gate/up MUL_MAT_ID fused with SwiGLU for compact expert tables. +template.decl @ggml.matmul_id.pair_f32_swiglu.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +amdgpu.target @ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_id_swiglu.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id_swiglu.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_swiglu.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_swiglu.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_swiglu.gate_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_swiglu.up_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def target(@ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32_gfx11_wave64) @ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index + %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_swiglu.input_size : index + %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_swiglu.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index + %gate_weight_format = config.get @ggml.mul_mat_id_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_id_swiglu.up_weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.pair_f32_swiglu.body>(%gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom new file mode 100644 index 000000000000..912fc0c9292c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_next_rmsnorm_publish_f32.loom @@ -0,0 +1,98 @@ +template.decl @ggml.matmul_id.tiled_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.tiled_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.tiled_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.mul_mat_id_postops.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_postops.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_postops.has_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.rms_epsilon : f32 + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.tiled_f32_f32.publish_vector4> device @ggml_mul_mat_id_tiled_postops_next_rmsnorm_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index + %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index + %c0 = index.constant 0 : index + %has_bias = index.cmp ne, %has_bias_value, %c0 : index + %has_residual = index.cmp ne, %has_residual_value, %c0 : index + scf.if %publish_word { + %with_bias = scf.if %has_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul_id.tiled_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %has_residual -> (vector<4xf32>) { + %added = template.apply<@ggml.matmul_id.tiled_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %added : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul_id.tiled_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %residual_output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul_id.tiled_f32_f32.finish_tile> device @ggml_mul_mat_id_tiled_postops_next_rmsnorm_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.apply<@ggml.matmul_id.tiled_f32_f32.finish_routed_rmsnorm_weight>(%token_count, %route_count, %output_size, %expert_count, %output_tile_count, %maximum_partition_count, %epsilon, %descriptor_ordinal, %expert, %route_tile_base, %partition_row_count, %expert_table, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, f32, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + template.return +} + +kernel.def target(@ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_postops.input_size : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_postops.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_id_postops.weight_format : index + %epsilon = config.get @ggml.mul_mat_id_postops.rms_epsilon : f32 + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.tiled_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %epsilon, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %residual_output, %residual_input, %residual_output, %norm_weight, %normalized_output, %completion_counters) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_publish_f32.loom new file mode 100644 index 000000000000..17b2a61a2c79 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_postops_publish_f32.loom @@ -0,0 +1,93 @@ +template.decl @ggml.matmul_id.tiled_f32_f32.add_bias_vector4(%arg0: index, %arg1: index, %arg2: vector<4xf32>, %arg3: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.tiled_f32_f32.add_residual_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) -> (vector<4xf32>) + +template.decl @ggml.matmul_id.tiled_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.store_vector4(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: vector<4xf32>, %arg6: buffer) + +amdgpu.target @ggml_mul_mat_id_tiled_input_f32_postops_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_postops.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id_postops.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id_postops.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_postops.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_postops.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_postops.has_bias : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_id_postops.has_residual : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.tiled_f32_f32.publish_vector4> device @ggml_mul_mat_id_tiled_postops_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %has_bias_value = config.get @ggml.mul_mat_id_postops.has_bias : index + %has_residual_value = config.get @ggml.mul_mat_id_postops.has_residual : index + %c0 = index.constant 0 : index + %has_bias = index.cmp ne, %has_bias_value, %c0 : index + %has_residual = index.cmp ne, %has_residual_value, %c0 : index + scf.if %publish_word { + %with_bias = scf.if %has_bias -> (vector<4xf32>) { + %biased = template.apply<@ggml.matmul_id.tiled_f32_f32.add_bias_vector4>(%output_size0, %channel, %values, %bias) : (index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %biased : vector<4xf32> + } else { + scf.yield %values : vector<4xf32> + } + %published = scf.if %has_residual -> (vector<4xf32>) { + %added = template.apply<@ggml.matmul_id.tiled_f32_f32.add_residual_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %with_bias, %residual_input) : (index, index, index, index, index, vector<4xf32>, buffer) -> (vector<4xf32>) + scf.yield %added : vector<4xf32> + } else { + scf.yield %with_bias : vector<4xf32> + } + template.apply<@ggml.matmul_id.tiled_f32_f32.store_vector4>(%token_count0, %route_count0, %output_size0, %assignment, %channel, %published, %output) : (index, index, index, index, index, vector<4xf32>, buffer) + } + template.return +} + +template.def<@ggml.matmul_id.tiled_f32_f32.finish_tile> device @ggml_mul_mat_id_tiled_postops_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_mul_mat_id_tiled_input_f32_postops_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_input_f32_postops_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %bias: buffer, %residual_input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_postops.input_size : index + %route_count = config.get @ggml.mul_mat_id_postops.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_postops.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_postops.expert_count : index + %output_size = config.get @ggml.mul_mat_id_postops.output_size : index + %weight_format = config.get @ggml.mul_mat_id_postops.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.tiled_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %bias, %output, %residual_input, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_publish_f32.loom new file mode 100644 index 000000000000..45a6e946eddb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_input_f32_publish_f32.loom @@ -0,0 +1,81 @@ +// Standalone GGML MUL_MAT_ID fallback for routed expert projections. +// +// The kernel consumes a precomputed expert table. Each table row stores compact +// assignment indices encoded as token * route_count + route. Input can either +// be broadcast across all routes or contain one input row per route. +template.decl @ggml.matmul_id.tiled_f32_f32.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: f32, %arg8: index, %arg9: index, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer, %arg18: buffer, %arg19: buffer, %arg20: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.matmul_id.tiled_f32_f32.publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +amdgpu.target @ggml_mul_mat_id_tiled_input_f32_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id.weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.matmul_id.tiled_f32_f32.publish_vector4> device @ggml_mul_mat_id_tiled_f32_f32_publish_vector4(%publish_word: i1, %assignment: index, %token: index, %route: index, %channel: index, %token_count0: index, %route_count0: index, %output_size0: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %route_count = index.assume %route_count0 [range(%route_count0, 1, 32)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %c0_offset = index.constant 0 : offset + %c1 = index.constant 1 : index + %output_row_count = index.mul %token_count, %route_count : index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%output_size]xf32> + scf.if %publish_word { + %bounded_assignment, %assignment_count = index.assume %assignment, %output_row_count [lt(%assignment, %output_row_count)] : index, index + %mask = vector.mask.range [%channel to %output_size step %c1] : index -> vector<4xi1> + vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<4xf32>, view<[%output_row_count]x[%output_size]xf32>, vector<4xi1> + } + template.return +} + +template.def<@ggml.matmul_id.tiled_f32_f32.finish_tile> device @ggml_mul_mat_id_tiled_f32_f32_finish_tile(%token_count: index, %route_count: index, %output_size: index, %expert_count: index, %output_tile_count: index, %maximum_partition_count: index, %epsilon: f32, %channel_tile: index, %descriptor_ordinal: index, %expert: index, %route_tile_base: index, %partition_row_count: index, %expert_table: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +kernel.def target(@ggml_mul_mat_id_tiled_input_f32_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_input_f32_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id.expert_count : index + %route_count = config.get @ggml.mul_mat_id.route_count : index + %output_size = config.get @ggml.mul_mat_id.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id.input_size : index + %route_count = config.get @ggml.mul_mat_id.route_count : index + %input_route_count = config.get @ggml.mul_mat_id.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id.expert_count : index + %output_size = config.get @ggml.mul_mat_id.output_size : index + %weight_format = config.get @ggml.mul_mat_id.weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.tiled_f32_f32.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %c0_f32, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom new file mode 100644 index 000000000000..3019636c626c --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/tiled_pair_input_f32_swiglu_publish_f32.loom @@ -0,0 +1,54 @@ +// Routed GGML gate/up MUL_MAT_ID fused with SwiGLU for tiled expert tables. +template.decl @ggml.matmul_id.pair_f32_swiglu.body(%gate_weight_format: index, %up_weight_format: index, %token_count: index, %input_size: index, %output_size: index, %route_count: index, %input_route_count: index, %expert_count: index, %channel_tile: index, %partition_ordinal: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) + +amdgpu.target @ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id_swiglu.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_id_swiglu.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_id_swiglu.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id_swiglu.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_swiglu.input_route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.mul_mat_id_swiglu.gate_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.mul_mat_id_swiglu.up_weight_format : %value: index where [range(%value, 4, 81)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def target(@ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32_gfx11_wave64) @ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32(%token_count: index) { + %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index + %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %assignment_count = index.mul %token_count, %route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index + %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index + kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %input_size = config.get @ggml.mul_mat_id_swiglu.input_size : index + %route_count = config.get @ggml.mul_mat_id_swiglu.route_count : index + %input_route_count = config.get @ggml.mul_mat_id_swiglu.input_route_count : index + %expert_count = config.get @ggml.mul_mat_id_swiglu.expert_count : index + %output_size = config.get @ggml.mul_mat_id_swiglu.output_size : index + %gate_weight_format = config.get @ggml.mul_mat_id_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_id_swiglu.up_weight_format : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %partition_ordinal = kernel.workgroup.id : index + template.apply<@ggml.matmul_id.pair_f32_swiglu.body>(%gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %route_count, %input_route_count, %expert_count, %channel_tile, %partition_ordinal, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_pair_binary.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_pair_binary.loom new file mode 100644 index 000000000000..f82513c00d0e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_pair_binary.loom @@ -0,0 +1,112 @@ +// Token-vector routed pair projection with compile-configured binary operation. +template.decl @ggml.matmul_id.vector_pair.body(%publish_output: i1, %binary_op: index, %input_size: index, %token_count: index, %token: index, %route_count: index, %route: index, %route_stride: index, %expert_count: index, %output_size: index, %channel: index, %lane: index, %input: buffer, %route_ids: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) + +template.decl @ggml.matmul_id.vector_pair.q8_1_x4.completion(%token_count: index, %route_count: index, %output_size: index, %token: index, %route: index, %channel_base: index, %subgroup: index, %workitem: index, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) + +config.decl @ggml.mul_mat_id.vector.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] + +config.decl @ggml.mul_mat_id.vector.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id.vector.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @ggml.mul_mat_id.vector.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @ggml.mul_mat_id.vector.input_format : %value: index where [range(%value, 0, 64)] + +config.decl @ggml.mul_mat_id.vector.lhs_weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.mul_mat_id.vector.rhs_weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.mul_mat_id.vector.binary_op : %value: index where [range(%value, 0, 8)] + +config.decl @ggml.mul_mat_id.vector.publish_format : %value: index where [range(%value, 0, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def @ggml_mul_mat_id_vector_pair_binary_publish_f32(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) { + %configured_route_count = config.get @ggml.mul_mat_id.vector.route_count : index + %configured_output_size = config.get @ggml.mul_mat_id.vector.output_size : index + %unit = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %workgroup_size = index.constant 128 : index + %padded_output_size = index.add %configured_output_size, %c3 : index + %channel_workgroup_count = index.div %padded_output_size, %c4 : index + kernel.launch.config workgroups(%channel_workgroup_count, %configured_route_count, %token_count) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_id.vector.input_size : index + %configured_route_count0 = config.get @ggml.mul_mat_id.vector.route_count : index + %configured_expert_count0 = config.get @ggml.mul_mat_id.vector.expert_count : index + %configured_output_size0 = config.get @ggml.mul_mat_id.vector.output_size : index + %input_format = config.get @ggml.mul_mat_id.vector.input_format : index + %lhs_weight_format = config.get @ggml.mul_mat_id.vector.lhs_weight_format : index + %rhs_weight_format = config.get @ggml.mul_mat_id.vector.rhs_weight_format : index + %binary_op = config.get @ggml.mul_mat_id.vector.binary_op : index + %publish_format = config.get @ggml.mul_mat_id.vector.publish_format : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512)] : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 128)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index + %token0 = kernel.workgroup.id : index + %channel_workgroup = kernel.workgroup.id : index + %route0 = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %safe_token0 = scf.select %valid_token, %token0, %c0 : index + %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index + %channel_base = index.mul %channel_workgroup, %c4 : index + %channel0 = index.add %channel_base, %subgroup : index + %valid_channel = index.cmp ult, %channel0, %bounded_output_size : index + %safe_channel0 = scf.select %valid_channel, %channel0, %c0 : index + %safe_channel, %body_output_size = index.assume %safe_channel0, %bounded_output_size [lt(%safe_channel0, %bounded_output_size)] : index, index + %route, %body_route_count = index.assume %route0, %bounded_route_count [lt(%route0, %bounded_route_count)] : index, index + %publishes_output = scalar.andi %valid_token, %valid_channel : i1 + template.apply<@ggml.matmul_id.vector_pair.body>(%publishes_output, %binary_op, %input_size, %body_token_count, %safe_token, %body_route_count, %route, %bounded_route_stride, %bounded_expert_count, %body_output_size, %safe_channel, %lane, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : (i1, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def @ggml_mul_mat_id_vector_pair_binary_publish_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) { + %configured_route_count = config.get @ggml.mul_mat_id.vector.route_count : index + %configured_output_size = config.get @ggml.mul_mat_id.vector.output_size : index + %unit = index.constant 1 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %workgroup_size = index.constant 128 : index + %padded_output_size = index.add %configured_output_size, %c3 : index + %channel_workgroup_count = index.div %padded_output_size, %c4 : index + kernel.launch.config workgroups(%channel_workgroup_count, %configured_route_count, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) { + %input_size = config.get @ggml.mul_mat_id.vector.input_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_route_count0 = config.get @ggml.mul_mat_id.vector.route_count : index + %configured_expert_count0 = config.get @ggml.mul_mat_id.vector.expert_count : index + %configured_output_size0 = config.get @ggml.mul_mat_id.vector.output_size : index + %input_format = config.get @ggml.mul_mat_id.vector.input_format : index + %lhs_weight_format = config.get @ggml.mul_mat_id.vector.lhs_weight_format : index + %rhs_weight_format = config.get @ggml.mul_mat_id.vector.rhs_weight_format : index + %binary_op = config.get @ggml.mul_mat_id.vector.binary_op : index + %publish_format = config.get @ggml.mul_mat_id.vector.publish_format : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 128)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %q8_input_noalias, %route_ids_noalias, %lhs_weight_noalias, %rhs_weight_noalias, %output_noalias, %completion_counters_noalias, %next_q8_output_noalias = buffer.assume.noalias %q8_input, %route_ids, %lhs_weight, %rhs_weight, %output, %completion_counters, %next_q8_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %publishes_binary = scalar.constant true : i1 + %body_token = index.constant 0 : index + %channel_workgroup = kernel.workgroup.id : index + %route = kernel.workgroup.id : index + %token = kernel.workgroup.id : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %workitem = kernel.workitem.id : index + %c4 = index.constant 4 : index + %channel_base = index.mul %channel_workgroup, %c4 : index + %channel = index.add %channel_base, %subgroup : index + template.apply<@ggml.matmul_id.vector_pair.body>(%publishes_binary, %binary_op, %input_size, %bounded_token_count, %body_token, %bounded_route_count, %route, %bounded_route_stride, %bounded_expert_count, %bounded_output_size, %channel, %lane, %q8_input_noalias, %route_ids_noalias, %lhs_weight_noalias, %rhs_weight_noalias, %output_noalias) : (i1, index, index, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.vector_pair.q8_1_x4.completion>(%bounded_token_count, %bounded_route_count, %bounded_output_size, %token, %route, %channel_base, %subgroup, %workitem, %output_noalias, %completion_counters_noalias, %next_q8_output_noalias) : (index, index, index, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_weighted.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_weighted.loom new file mode 100644 index 000000000000..c83addffba4e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/matmul_id/vector_weighted.loom @@ -0,0 +1,88 @@ +// Token-vector routed weighted projection with next Q8_1 x4 publication. +template.decl @ggml.matmul_id.vector_weighted.q4k_q8.wave32_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul_id.vector_weighted.q6k_f32.wave64_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.matmul_id.vector_weighted.next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) + +amdgpu.target @ggml_mul_mat_id_vector_weighted_wave32_publish_q8_gfx11 {subgroup_size = 32} + +amdgpu.target @ggml_mul_mat_id_vector_weighted_wave64_publish_q8_gfx11 {subgroup_size = 64} + +config.decl @ggml.mul_mat_id.vector.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_id.vector.route_count : %value: index where [range(%value, 1, 8)] + +config.decl @ggml.mul_mat_id.vector.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.mul_mat_id.vector.output_size : %value: index where [range(%value, 1, 4096)] + +config.decl @ggml.mul_mat_id.vector.input_format : %value: index where [range(%value, 0, 64)] + +config.decl @ggml.mul_mat_id.vector.weight_format : %value: index where [range(%value, 0, 81)] + +config.decl @ggml.mul_mat_id.vector.publish_format : %value: index where [range(%value, 0, 64)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def target(@ggml_mul_mat_id_vector_weighted_wave32_publish_q8_gfx11) @ggml_mul_mat_id_vector_weighted_wave32_publish_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %configured_output_size = config.get @ggml.mul_mat_id.vector.output_size : index + %c8 = index.constant 8 : index + %c7 = index.constant 7 : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_input_size0 = config.get @ggml.mul_mat_id.vector.input_size : index + %configured_route_count0 = config.get @ggml.mul_mat_id.vector.route_count : index + %configured_expert_count0 = config.get @ggml.mul_mat_id.vector.expert_count : index + %configured_output_size0 = config.get @ggml.mul_mat_id.vector.output_size : index + %input_format = config.get @ggml.mul_mat_id.vector.input_format : index + %weight_format = config.get @ggml.mul_mat_id.vector.weight_format : index + %publish_format = config.get @ggml.mul_mat_id.vector.publish_format : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %c8 = index.constant 8 : index + %publishes_output = scalar.constant true : i1 + %body_token = index.constant 0 : index + template.apply<@ggml.matmul_id.vector_weighted.q4k_q8.wave32_body>(%publishes_output, %bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.vector_weighted.next_q8_completion>(%c8, %c8, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_mul_mat_id_vector_weighted_wave64_publish_q8_gfx11) @ggml_mul_mat_id_vector_weighted_wave64_publish_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { + %configured_output_size = config.get @ggml.mul_mat_id.vector.output_size : index + %c1 = index.constant 1 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %configured_output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %configured_input_size0 = config.get @ggml.mul_mat_id.vector.input_size : index + %configured_route_count0 = config.get @ggml.mul_mat_id.vector.route_count : index + %configured_expert_count0 = config.get @ggml.mul_mat_id.vector.expert_count : index + %configured_output_size0 = config.get @ggml.mul_mat_id.vector.output_size : index + %input_format = config.get @ggml.mul_mat_id.vector.input_format : index + %weight_format = config.get @ggml.mul_mat_id.vector.weight_format : index + %publish_format = config.get @ggml.mul_mat_id.vector.publish_format : index + %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %body_token = index.constant 0 : index + template.apply<@ggml.matmul_id.vector_weighted.q6k_f32.wave64_pair_body>(%bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %input, %route_ids, %route_weights, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + template.apply<@ggml.matmul_id.vector_weighted.next_q8_completion>(%c8, %c4, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/moe_routing_tables.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/moe_routing_tables.loom new file mode 100644 index 000000000000..525081faef51 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/moe_routing_tables.loom @@ -0,0 +1,178 @@ +// Common MoE routing table preparation. +// +// The expert table stores [expert_count] assignment counts followed by +// [expert_count][token_count] compact assignment ordinals. The partition table +// stores a descriptor count followed by compact 32-row expert partitions. +amdgpu.target @ggml_moe_routing_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.moe_routing.route_count : %value: index where [range(%value, 1, 32)] + +config.decl @ggml.moe_routing.expert_count : %value: index where [range(%value, 1, 512)] + +config.decl @ggml.moe_routing.descriptor_expert_mask : %value: i32 + +config.decl @ggml.moe_routing.descriptor_partition_shift : %value: i32 + +config.decl @ggml.moe_routing.descriptor_row_count_shift : %value: i32 + +config.decl @ggml.moe_routing.partition_workgroup_size : %value: index where [range(%value, 1, 1024)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def target(@ggml_moe_routing_gfx11_wave32) @ggml_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %configured_expert_count = config.get @ggml.moe_routing.expert_count : index + %c1 = index.constant 1 : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%configured_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) where [range(%token_count, 1, 2048)] { + %configured_route_count0 = config.get @ggml.moe_routing.route_count : index + %configured_expert_count0 = config.get @ggml.moe_routing.expert_count : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 512), ge(%route_stride, %bounded_route_count)] : index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %expert0 = kernel.workgroup.id : index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %workgroup_size = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %cn1_i32 = scalar.constant -1 : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %expert, %table_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index + %assignment_table_byte_base = index.scale %table_expert_count, %c4_bytes : index, offset -> offset + %route_view = buffer.view %route_ids[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %count_view = buffer.view %expert_table[%c0_offset] : buffer -> view<[%table_expert_count]xi32> + %assignment_view = buffer.view %expert_table[%assignment_table_byte_base] : buffer -> view<[%table_expert_count]x[%bounded_token_count]xi32> + %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { + %assignment = index.add %block_base, %lane : index + %in_range = index.cmp ult, %assignment, %assignment_count : index + %route_expert_i32 = scf.if %in_range -> (i32) { + %token0 = index.div %assignment, %configured_route_count : index + %route0 = index.rem %assignment, %configured_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %cn1_i32 : i32 + } + %route_expert0 = index.cast %route_expert_i32 : i32 to index + %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 511)] : index + %matches = index.cmp eq, %route_expert, %expert : index + %match_i32 = scf.if %matches -> (i32) { + scf.yield %c1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 + %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 + %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 + scf.if %matches { + %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 + %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index + %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 2047)] : index + %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%table_expert_count]x[%bounded_token_count]xi32> + } + %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 + scf.yield %next_matched_base : i32 + } + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %expert_route_count, %count_view[%expert] : i32, view<[%table_expert_count]xi32> + } + kernel.return +} + +kernel.def target(@ggml_moe_routing_gfx11_wave32) @ggml_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) { + %c1 = index.constant 1 : index + %workgroup_size = config.get @ggml.moe_routing.partition_workgroup_size : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) where [range(%token_count, 1, 2048)] { + %configured_route_count0 = config.get @ggml.moe_routing.route_count : index + %configured_expert_count0 = config.get @ggml.moe_routing.expert_count : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index + %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index + %lane = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %descriptor_expert_mask = config.get @ggml.moe_routing.descriptor_expert_mask : i32 + %descriptor_partition_shift = config.get @ggml.moe_routing.descriptor_partition_shift : i32 + %descriptor_row_count_shift = config.get @ggml.moe_routing.descriptor_row_count_shift : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index + %count_view = buffer.view %expert_table[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> + %partition_count_view = buffer.view %partition_table[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %has_expert = index.cmp ult, %lane, %bounded_expert_count : index + %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { + %expert, %table_expert_count = index.assume %lane, %bounded_expert_count [lt(%lane, %bounded_expert_count)] : index, index + %loaded = view.load %count_view[%expert] : view<[%bounded_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index + %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 2048)] : index + %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index + %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index + %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 + %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 + %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 + %partition_count0 = index.cast %partition_count_i32 : i32 to index + %partition_count, %table_partition_capacity = index.assume %partition_count0, %maximum_partition_count [le(%partition_count0, %maximum_partition_count)] : index, index + %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index + %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 2559)] : index + scf.if %has_expert { + %expert, %table_expert_count = index.assume %lane, %bounded_expert_count [lt(%lane, %bounded_expert_count)] : index, index + %expert_i32 = index.cast %expert : index to i32 + scf.for %partition = [%c0 to %expert_partition_count step %c1] { + %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index + %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %partition_count [lt(%descriptor_ordinal0, %partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %partition_remainder = index.rem %expert_assignment_count, %c32 : index + %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index + %partition_row_count = scf.if %has_partial_tail -> (index) { + %next_partition = index.add %partition, %c1 : index + %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index + %tail_row_count = scf.if %is_tail_partition -> (index) { + scf.yield %partition_remainder : index + } else { + scf.yield %c32 : index + } + scf.yield %tail_row_count : index + } else { + scf.yield %c32 : index + } + %partition_i32 = index.cast %partition : index to i32 + %partition_row_count_i32 = index.cast %partition_row_count : index to i32 + %bounded_expert_i32 = scalar.andi %expert_i32, %descriptor_expert_mask : i32 + %packed_partition = scalar.shli %partition_i32, %descriptor_partition_shift : i32 + %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 + %packed_row_count = scalar.shli %partition_row_count_minus_one, %descriptor_row_count_shift : i32 + %packed_expert_partition = scalar.ori %bounded_expert_i32, %packed_partition : i32 + %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 + view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> + } + } + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_decode.loom new file mode 100644 index 000000000000..22e7cc97c168 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_add_f32_f32_decode.loom @@ -0,0 +1,55 @@ +template.decl @ggml.mul_mat_f32_f32_decode.dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output: buffer, %residual_input: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) + +amdgpu.target @ggml_mul_mat_add_f32_f32_decode_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_f32_f32_decode.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.mul_mat_f32_f32_decode.output_capacity : %value: index where [range(%value, 1, 1048576)] + +config.decl @ggml.mul_mat_f32_f32_decode.weight_format : %value: index where [range(%value, 4, 81)] + +template.def<@ggml.mul_mat_f32_f32_decode.publish_pair2> device priority(20) @ggml_mul_mat_add_f32_f32_decode_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output: buffer, %residual_input: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 1048576)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %residual_input_noalias = buffer.assume.noalias %residual_input : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%output_size]xf32> + %residual_input_view = buffer.view %residual_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%output_size]xf32> + scf.if %publish_pair { + %residual0 = view.load %residual_input_view[%token, %channel] : view<[%launch_token_count]x[%output_size]xf32> -> f32 + %sum0 = scalar.addf %value0, %residual0 : f32 + view.store %sum0, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %residual1 = view.load %residual_input_view[%token, %row1] : view<[%launch_token_count]x[%output_size]xf32> -> f32 + %sum1 = scalar.addf %value1, %residual1 : f32 + view.store %sum1, %output_view[%token, %row1] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + } + template.return +} + +kernel.def target(@ggml_mul_mat_add_f32_f32_decode_gfx11_wave64) @ggml_mul_mat_add_f32_f32_decode_wave64(%token_count: index, %input_size: index, %output_size: index) { + %token_capacity = config.get @ggml.mul_mat_f32_f32_decode.token_capacity : index + %output_capacity = config.get @ggml.mul_mat_f32_f32_decode.output_capacity : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %padded_output_size = index.add %output_capacity, %c1 : index + %output_pairs = index.div %padded_output_size, %c2 : index + %padded_token_count = index.add %token_capacity, %c63 : index + %launch_tokens = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_pairs, %launch_tokens, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %residual_input: buffer, %residual_output: buffer) { + %weight_format = config.get @ggml.mul_mat_f32_f32_decode.weight_format : index + template.apply<@ggml.mul_mat_f32_f32_decode.dispatch>(%weight_format, %token_count, %input_size, %output_size, %input, %weight, %residual_output, %residual_input, %residual_output, %residual_output, %residual_output) : (index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_dual_q4_f32_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_dual_q4_f32_decode.loom new file mode 100644 index 000000000000..8bbc43af5d39 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_dual_q4_f32_decode.loom @@ -0,0 +1,88 @@ +// Two raw-Q4_K one-token F32 projections sharing activation loads. +// Each wave64 computes the same row of two distinct weights. The original +// public decoder retains F16 weight rounding, four strided K cohorts, dot4 +// contraction and reduction order. Both ordinary F32 outputs are published. +config.decl @ggml.mul_mat_dual_q4_f32_c1.input_size : %value: index where [range(%value, 4096, 32768), mul(%value, 1024)] +config.decl @ggml.mul_mat_dual_q4_f32_c1.output_size : %value: index where [range(%value, 24, 63)] + +amdgpu.target @ggml_mul_mat_dual_q4_f32_decode_gfx11_wave64 {subgroup_size = 64} + +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_mul_mat_decode_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>) + +func.decl @ggml_mul_mat_decode_f32_block_row(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %input_size: index, %row: index, %block: index, %lane: index, %weight: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) + +func.def inline @ggml_mul_mat_dual_q4_f32_decode_body(%weight_format: index, %publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 32)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 1048576)] : 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 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.add %pair, %c0 : index + %row0, %launch_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index + %row1 = index.add %row0, %c0 : index + %row1_valid = index.cmp ult, %row1, %launch_output_size : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %c0_f32 : f32, %row_acc1 = %c0_f32 : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + %input0, %input1, %input2, %input3 = func.call @ggml_mul_mat_decode_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %bounded_input_size, %row0, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + %row1_contribution = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %bounded_input_size, %row1, %block, %lane, %aux0, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %c0_f32 : 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 + %alpha_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_output_size]xf32> + %beta_view = buffer.view %aux1[%c0_offset] : buffer -> view<[%launch_output_size]xf32> + scf.if %publish_pair { + view.store %sum0, %alpha_view[%row0] : f32, view<[%launch_output_size]xf32> + view.store %sum1, %beta_view[%row0] : f32, view<[%launch_output_size]xf32> + } + } + func.return +} + +kernel.def target(@ggml_mul_mat_dual_q4_f32_decode_gfx11_wave64) export("ggml_mul_mat_dual_q4_f32_decode") @ggml_mul_mat_dual_q4_f32_decode() { + %heads = config.get @ggml.mul_mat_dual_q4_f32_c1.output_size : index + %threads = index.constant 64 : index + %one = index.constant 1 : index + kernel.launch.config workgroups(%heads, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%input: buffer, %alpha: buffer, %beta: buffer, %output_alpha: buffer, %output_beta: buffer) { + %format = index.constant 4 : index + %one = index.constant 1 : index + %zero = index.constant 0 : index + %yes = scalar.constant true : i1 + %k = config.get @ggml.mul_mat_dual_q4_f32_c1.input_size : index + %n = config.get @ggml.mul_mat_dual_q4_f32_c1.output_size : index + %row = kernel.workgroup.id : index + func.call @ggml_mul_mat_dual_q4_f32_decode_body(%format, %yes, %one, %zero, %row, %k, %n, %input, %alpha, %output_alpha, %beta, %output_beta, %output_alpha, %output_alpha) : (index, i1, 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/mul_mat_f32_f32_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom new file mode 100644 index 000000000000..71a1996fd358 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_decode.loom @@ -0,0 +1,400 @@ +// Dense GGML decode matmul for one contiguous F32 activation row. +// +// Pipeline: +// 1. Load four masked F32 activation packets from each 256-wide K block. +// 2. Decode four matching weight packets through @ggml_dequant_f16_vector4. +// 3. Dot the activation/weight packets and reduce one contribution per output row. +// +// One 64-workitem workgroup computes two adjacent output rows. Four cohorts +// of 16 lanes each process four 256-value weight blocks in parallel, while +// each lane contracts four packets from the 0, 32, 64, and 96 element +// quarters of its block. Storage-format-specific weight decoding stays in +// dequant.loom, so decode scheduling remains format-agnostic. +template.decl @ggml.mul_mat_f32_f32_decode.body(%weight_format: index, %publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.launch(%token_capacity: index, %output_capacity: index) -> (index, index, index, index) + +template.decl @ggml.mul_mat_f32_f32_decode.publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.dual_body(%op: index, %lhs_weight_format: index, %rhs_weight_format: index, %publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.dual_dispatch(%op: index, %lhs_weight_format: index, %rhs_weight_format: index, %token_capacity: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.publish_dual_pair2(%op: index, %publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %lhs0: f32, %lhs1: f32, %rhs0: f32, %rhs1: f32, %output: buffer) + +amdgpu.target @ggml_mul_mat_f32_f32_decode_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.mul_mat_f32_f32_decode.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.mul_mat_f32_f32_decode.output_capacity : %value: index where [range(%value, 1, 1048576)] + +config.decl @ggml.mul_mat_f32_f32_decode.weight_format : %value: index where [range(%value, 4, 81)] + +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_mul_mat_decode_dot4_f32(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (f32) { + %c0 = scalar.constant 0.0 : f32 + %lhs0 = vector.extract %lhs[0] : vector<4xf32> -> f32 + %lhs1 = vector.extract %lhs[1] : vector<4xf32> -> f32 + %lhs2 = vector.extract %lhs[2] : vector<4xf32> -> f32 + %lhs3 = vector.extract %lhs[3] : vector<4xf32> -> f32 + %rhs0 = vector.extract %rhs[0] : vector<4xf32> -> f32 + %rhs1 = vector.extract %rhs[1] : vector<4xf32> -> f32 + %rhs2 = vector.extract %rhs[2] : vector<4xf32> -> f32 + %rhs3 = vector.extract %rhs[3] : vector<4xf32> -> f32 + %sum0 = scalar.fmaf %lhs0, %rhs0, %c0 : f32 + %sum1 = scalar.fmaf %lhs1, %rhs1, %sum0 : f32 + %sum2 = scalar.fmaf %lhs2, %rhs2, %sum1 : f32 + %sum3 = scalar.fmaf %lhs3, %rhs3, %sum2 : f32 + func.return %sum3 : f32 +} + +// Load activation packets once per K block. The masks zero K-tail lanes for +// input sizes that are not exact 256 multiples. +func.def inline @ggml_mul_mat_decode_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 + %c0_f32x4 = 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, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input1 = vector.load.mask %input_view[%token, %safe_block, %input_index1], %mask1, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input2 = vector.load.mask %input_view[%token, %safe_block, %input_index2], %mask2, %c0_f32x4 : view<[%token_count]x[%launch_block_count]x256xf32>, vector<4xi1>, vector<4xf32> + %input3 = vector.load.mask %input_view[%token, %safe_block, %input_index3], %mask3, %c0_f32x4 : 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> +} + +// Decode one output-row contribution for the current K block. The four weight +// packets mirror the activation packet offsets: 0, 32, 64, and 96. +func.def inline @ggml_mul_mat_decode_f32_block_row(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %input_size: index, %row: index, %block: index, %lane: index, %weight: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c96 = index.constant 96 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_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 + %weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_byte_base = index.scale %row, %weight_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 + %vector_half_base = index.mul %vector_half, %c128 : index + %packet_k = index.mul %packet, %c4 : index + %k_half = index.add %block_k_base, %vector_half_base : index + %k0 = index.add %k_half, %packet_k : index + %k1_add = index.add %k0, %c32 : index + %k2_add = index.add %k0, %c64 : index + %k3_add = index.add %k0, %c96 : index + %valid_k0 = index.cmp ult, %k0, %input_size : index + %valid_k1 = index.cmp ult, %k1_add, %input_size : index + %valid_k2 = index.cmp ult, %k2_add, %input_size : index + %valid_k3 = index.cmp ult, %k3_add, %input_size : index + %weight0_f16 = scf.if %valid_k0 -> (vector<4xf16>) { + %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, %row_byte_base, %input_size, %safe_block, %group0, %packet, %k0) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %weight1_f16 = scf.if %valid_k1 -> (vector<4xf16>) { + %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, %row_byte_base, %input_size, %safe_block, %group1, %packet, %k1_add) : (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> + } + %weight2_f16 = scf.if %valid_k2 -> (vector<4xf16>) { + %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, %row_byte_base, %input_size, %safe_block, %group2, %packet, %k2_add) : (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> + } + %weight3_f16 = scf.if %valid_k3 -> (vector<4xf16>) { + %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, %row_byte_base, %input_size, %safe_block, %group3, %packet, %k3_add) : (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> + } + %weight0 = vector.extf %weight0_f16 : vector<4xf16> to vector<4xf32> + %weight1 = vector.extf %weight1_f16 : vector<4xf16> to vector<4xf32> + %weight2 = vector.extf %weight2_f16 : vector<4xf16> to vector<4xf32> + %weight3 = vector.extf %weight3_f16 : vector<4xf16> to vector<4xf32> + %dot0 = func.call @ggml_mul_mat_decode_dot4_f32(%input0, %weight0) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot1 = func.call @ggml_mul_mat_decode_dot4_f32(%input1, %weight1) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot2 = func.call @ggml_mul_mat_decode_dot4_f32(%input2, %weight2) : (vector<4xf32>, vector<4xf32>) -> (f32) + %dot3 = func.call @ggml_mul_mat_decode_dot4_f32(%input3, %weight3) : (vector<4xf32>, vector<4xf32>) -> (f32) + %sum3 = scalar.addf %dot2, %dot3 : f32 + %sum2 = scalar.addf %dot1, %sum3 : f32 + %sum = scalar.addf %dot0, %sum2 : f32 + %contribution = scf.select %valid_block, %sum, %c0_f32 : f32 + func.return %contribution : f32 +} + +template.def<@ggml.mul_mat_f32_f32_decode.publish_pair2> device priority(1) @ggml_mul_mat_f32_f32_decode_publish_pair2(%publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %value0: f32, %value1: f32, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 1048576)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%output_size]xf32> + scf.if %publish_pair { + view.store %value0, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + view.store %value1, %output_view[%token, %row1] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_decode.publish_dual_pair2> device priority(1) @ggml_mul_mat_f32_f32_decode_publish_dual_pair2(%op: index, %publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %lhs0: f32, %lhs1: f32, %rhs0: f32, %rhs1: f32, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 1048576)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%output_size]xf32> + scf.if %publish_pair { + view.store %lhs0, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + view.store %lhs1, %output_view[%token, %row1] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_decode.body> device @ggml_mul_mat_f32_f32_decode_body(%weight_format: index, %publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 32)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 1048576)] : 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 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.mul %pair, %c2 : index + %row0, %launch_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %launch_output_size : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %c0_f32 : f32, %row_acc1 = %c0_f32 : f32) -> (f32, f32) { + %block = index.add %block_base, %cohort : index + %input0, %input1, %input2, %input3 = func.call @ggml_mul_mat_decode_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %contribution0 = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %bounded_input_size, %row0, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %contribution1 = scf.if %row1_valid -> (f32) { + %row1_contribution = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %bounded_input_size, %row1, %block, %lane, %weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %row1_contribution : f32 + } else { + scf.yield %c0_f32 : 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.mul_mat_f32_f32_decode.publish_pair2>(%publish_pair, %token, %row0, %launch_token_count, %launch_output_size, %row1_valid, %sum0, %sum1, %output_noalias, %aux0, %aux1, %aux2, %aux3) : (i1, index, index, index, index, i1, f32, f32, buffer, buffer, buffer, buffer, buffer) + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_decode.dual_body> device @ggml_mul_mat_f32_f32_decode_dual_body(%op: index, %lhs_weight_format: index, %rhs_weight_format: index, %publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index + %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 32)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 1048576)] : 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 + %c0_f32 = scalar.constant 0.0 : f32 + %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %row00 = index.mul %pair, %c2 : index + %row0, %launch_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index + %row1 = index.add %row0, %c1 : index + %row1_valid = index.cmp ult, %row1, %launch_output_size : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %block_count = index.div %padded_input_size, %c256 : index + %cohort = index.div %lane, %c16 : index + scf.if %publish_output { + %input_noalias, %lhs_weight_noalias, %rhs_weight_noalias, %output_noalias = buffer.assume.noalias %input, %lhs_weight, %rhs_weight, %output : buffer, buffer, buffer, buffer + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %lhs_acc0, %lhs_acc1, %rhs_acc0, %rhs_acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%lhs_row_acc0 = %c0_f32 : f32, %lhs_row_acc1 = %c0_f32 : f32, %rhs_row_acc0 = %c0_f32 : f32, %rhs_row_acc1 = %c0_f32 : f32) -> (f32, f32, f32, f32) { + %block = index.add %block_base, %cohort : index + %input0, %input1, %input2, %input3 = func.call @ggml_mul_mat_decode_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) + %lhs_contribution0 = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %lhs_weight_format, %bounded_input_size, %row0, %block, %lane, %lhs_weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %rhs_contribution0 = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %rhs_weight_format, %bounded_input_size, %row0, %block, %lane, %rhs_weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %lhs_contribution1, %rhs_contribution1 = scf.if %row1_valid -> (f32, f32) { + %lhs_row1_contribution = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %lhs_weight_format, %bounded_input_size, %row1, %block, %lane, %lhs_weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + %rhs_row1_contribution = func.call @ggml_mul_mat_decode_f32_block_row(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %rhs_weight_format, %bounded_input_size, %row1, %block, %lane, %rhs_weight_noalias, %input0, %input1, %input2, %input3) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, index, index, index, index, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) + scf.yield %lhs_row1_contribution, %rhs_row1_contribution : f32, f32 + } else { + scf.yield %c0_f32, %c0_f32 : f32, f32 + } + %lhs_next0 = scalar.addf %lhs_row_acc0, %lhs_contribution0 : f32 + %lhs_next1 = scalar.addf %lhs_row_acc1, %lhs_contribution1 : f32 + %rhs_next0 = scalar.addf %rhs_row_acc0, %rhs_contribution0 : f32 + %rhs_next1 = scalar.addf %rhs_row_acc1, %rhs_contribution1 : f32 + scf.yield %lhs_next0, %lhs_next1, %rhs_next0, %rhs_next1 : f32, f32, f32, f32 + } + %lhs_sum0 = kernel.workgroup.reduce %lhs_acc0 : f32 + %lhs_sum1 = kernel.workgroup.reduce %lhs_acc1 : f32 + %rhs_sum0 = kernel.workgroup.reduce %rhs_acc0 : f32 + %rhs_sum1 = kernel.workgroup.reduce %rhs_acc1 : f32 + %is_lane_zero = index.cmp eq, %lane, %c0 : index + %publish_pair = scalar.andi %publish_output, %is_lane_zero : i1 + template.apply<@ggml.mul_mat_f32_f32_decode.publish_dual_pair2>(%op, %publish_pair, %token, %row0, %launch_token_count, %launch_output_size, %row1_valid, %lhs_sum0, %lhs_sum1, %rhs_sum0, %rhs_sum1, %output_noalias) : (index, i1, index, index, index, index, i1, f32, f32, f32, f32, buffer) + } + template.return +} + +template.def<@ggml.mul_mat_f32_f32_decode.launch> @ggml_mul_mat_f32_f32_decode_launch(%token_capacity: index, %output_capacity: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c64 = index.constant 64 : index + %padded_output_size = index.add %output_capacity, %c1 : index + %output_pairs = index.div %padded_output_size, %c2 : index + template.return %output_pairs, %token_capacity, %c1, %c64 : index, index, index, index +} + +template.def<@ggml.mul_mat_f32_f32_decode.dual_dispatch> device @ggml_mul_mat_f32_f32_decode_dual_dispatch(%op: index, %lhs_weight_format: index, %rhs_weight_format: index, %token_capacity: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 1048576)] : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %pair = kernel.workgroup.id : index + %token = kernel.workgroup.id : index + %row0 = index.mul %pair, %c2 : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_row = index.cmp ult, %row0, %bounded_output_size : index + %publish_output = scalar.andi %valid_token, %valid_row : i1 + %safe_token = scf.select %valid_token, %token, %c0 : index + %safe_pair = scf.select %valid_row, %pair, %c0 : index + template.apply<@ggml.mul_mat_f32_f32_decode.dual_body>(%op, %lhs_weight_format, %rhs_weight_format, %publish_output, %bounded_token_count, %safe_token, %safe_pair, %input_size, %bounded_output_size, %input, %lhs_weight, %rhs_weight, %output) : (index, index, index, i1, index, index, index, index, index, buffer, buffer, buffer, buffer) + template.return +} + +template.def<@ggml.mul_mat_f32_f32_decode.dispatch> device @ggml_mul_mat_f32_f32_decode_dispatch(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer, %aux0: buffer, %aux1: buffer, %aux2: buffer, %aux3: buffer) { + %token_capacity = config.get @ggml.mul_mat_f32_f32_decode.token_capacity : index + %output_capacity = config.get @ggml.mul_mat_f32_f32_decode.output_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 1, 1048576), le(%output_size, %output_capacity)] : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %pair = kernel.workgroup.id : index + %token = kernel.workgroup.id : index + %row0 = index.mul %pair, %c2 : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_row = index.cmp ult, %row0, %bounded_output_size : index + %publish_output = scalar.andi %valid_token, %valid_row : i1 + %safe_token = scf.select %valid_token, %token, %c0 : index + %safe_pair = scf.select %valid_row, %pair, %c0 : index + template.apply<@ggml.mul_mat_f32_f32_decode.body>(%weight_format, %publish_output, %bounded_token_count, %safe_token, %safe_pair, %input_size, %bounded_output_size, %input, %weight, %output, %aux0, %aux1, %aux2, %aux3) : (index, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.return +} + +kernel.def target(@ggml_mul_mat_f32_f32_decode_gfx11_wave64) @ggml_mul_mat_f32_f32_decode_wave64(%token_count: index, %input_size: index, %output_size: index) { + %token_capacity = config.get @ggml.mul_mat_f32_f32_decode.token_capacity : index + %output_capacity = config.get @ggml.mul_mat_f32_f32_decode.output_capacity : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %padded_output_size = index.add %output_capacity, %c1 : index + %output_pairs = index.div %padded_output_size, %c2 : index + %padded_token_count = index.add %token_capacity, %c63 : index + %launch_tokens = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_pairs, %launch_tokens, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { + %weight_format = config.get @ggml.mul_mat_f32_f32_decode.weight_format : index + template.apply<@ggml.mul_mat_f32_f32_decode.dispatch>(%weight_format, %token_count, %input_size, %output_size, %input, %weight, %output, %output, %output, %output, %output) : (index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom new file mode 100644 index 000000000000..6e7f8899cdfb --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom @@ -0,0 +1,217 @@ +func.decl @ggml_mul_mat_quantized_f16_prefill_conv4_interior(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %input: buffer, %weight: buffer, %filter: buffer, %output: buffer, %edges: buffer) + +func.decl @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer) + +// Dense GGML matmul for contiguous 2D F32 activations. +// +// Each two-wave workgroup computes 64 output channels for 32 contiguous +// tokens. GGUF weight rows are decoded or converted directly into a padded +// FP16 LDS tile, while the same workgroup converts the corresponding F32 +// activation tile to FP16. Four wave64 WMMA accumulators cover the 32x32 result +// owned by each wave. A wave-private LDS slice transposes each accumulator for +// coalesced F32 publication. +// +// Weight contracts use the unmodified GGUF layout: +// Q3_K: [output channel][input size / 256][110 bytes] +// Q4_K: [output channel][input size / 256][144 bytes] +// Q5_K: [output channel][input size / 256][176 bytes] +// Q6_K: [output channel][input size / 256][210 bytes] +// IQ4_NL: [output channel][input size / 32][18 bytes] +// IQ3_S: [output channel][input size / 256][110 bytes] +// IQ4_XS: [output channel][input size / 256][136 bytes] +// Q8_0: [output channel][input size / 32][34 bytes] +// Q8_1: [output channel][input size / 32][36 bytes] +// F16: [output channel][input size][2 bytes] +// BF16: [output channel][input size][2 bytes] +// F32: [output channel][input size][4 bytes] +// The configured weight format selects the storage reader before the shared +// device template is instantiated, so inactive decode logic is absent from the +// emitted kernel. No persistent repacking or expanded-weight allocation is +// required. +template.decl @ggml.mul_mat_f32_f32_wmma.core(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: f32, %arg7: index, %arg8: index, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer, %arg17: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.mul_mat_f32_f32_wmma.launch(%token_capacity: index, %input_size: index, %output_size: index, %weight_format: index) -> (index, index, index, index) + +template.decl @ggml.mul_mat_f32_f32_wmma.publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) + +template.decl @ggml.unary_f32.apply_vector4(%arg0: index, %arg1: vector<4xf32>) -> (vector<4xf32>) + +amdgpu.target @ggml_mul_mat_gfx11_wave64 {subgroup_size = 64} +amdgpu.target @ggml_mul_mat_gfx11_wave32 {subgroup_size = 32} + +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)] + +// Private F16 activation storage for the Q4 prefill entry: row-major (0) or +// K16-major [K/16, tokens, 16] (1). The latter uses full 512-token tiles. +config.decl @ggml.mul_mat.f16_input_layout : %value: index where [range(%value, 0, 1)] + +// Selects whether publication overwrites the output (0) or accumulates the +// projection into a caller-provided residual (1). +config.decl @ggml.mul_mat.output_accumulation : %value: index where [range(%value, 0, 1)] + +// Unary op applied to the published F32 output. Identity is op 23. +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.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +template.def<@ggml.mul_mat_f32_f32_wmma.publish_vector4> device @ggml_mul_mat_f32_f32_wmma_publish_vector4(%publish_word: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %output_accumulation: index, %output_unary_op: index, %values: vector<4xf32>, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + %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 + %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.return +} + +template.def<@ggml.mul_mat_f32_f32_wmma.finish_tile> device @ggml_mul_mat_f32_f32_wmma_finish_tile(%token_count: index, %output_size: index, %output_tile_count: index, %token_tile_count: index, %output_accumulation: index, %output_unary_op: index, %epsilon: f32, %channel_tile: index, %token_tile: index, %bias: buffer, %output: buffer, %residual_input: buffer, %residual_output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counters: buffer) { + template.return +} + +// Shape and format configs specialize the native code once while token count +// remains the only per-dispatch scalar. +kernel.def target(@ggml_mul_mat_gfx11_wave64) @ggml_mul_mat_f32_f32_wmma(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %input_size = config.get @ggml.mul_mat.input_size : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %launch_x, %launch_y, %launch_z, %workgroup_size = template.apply<@ggml.mul_mat_f32_f32_wmma.launch>(%token_capacity, %input_size, %output_size, %weight_format) pure : (index, index, index, index) -> (index, index, index, index) + kernel.launch.config workgroups(%launch_x, %launch_y, %launch_z) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %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 + %c0_f32 = scalar.constant 0.0 : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_f32_f32_wmma.core>(%weight_format, %bounded_token_count, %input_size, %output_size, %output_accumulation, %output_unary_op, %c0_f32, %channel_tile, %token_tile, %input, %weight, %output, %output, %output, %output, %output, %output, %output) : (index, index, index, index, index, index, f32, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_mul_mat_gfx11_wave32) @ggml_mul_mat_q4_k_f16_wmma_prefill_wave32(%token_count: index) { + %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 + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c512 = index.constant 512 : index + %c1024 = index.constant 1024 : index + %minimum_workgroups = index.constant 32 : index + %twice_input_size = index.mul %input_size, %c2 : index + %contracts = index.cmp ule, %output_size, %input_size : index + %large_expansion = index.cmp uge, %output_size, %twice_input_size : index + %wide_shape = scalar.ori %contracts, %large_expansion : i1 + %output_tail = index.rem %output_size, %c128 : index + %full_wide_tile = index.cmp eq, %output_tail, %c0 : index + %wide_tile = scalar.andi %wide_shape, %full_wide_tile : i1 + %ordinary_tile_channels = scf.select %wide_tile, %c128, %c64 : index + %ordinary_workgroup_size = scf.select %wide_tile, %c512, %c256 : index + %active_token_tiles = index.div %token_count, %c256 : index + %larger_output_tiles = index.div %output_size, %c256 : index + %larger_grid = index.mul %active_token_tiles, %larger_output_tiles : index + %enough_workgroups = index.cmp uge, %larger_grid, %minimum_workgroups : index + %larger_channel_tail = index.rem %output_size, %c256 : index + %full_larger_channels = index.cmp eq, %larger_channel_tail, %c0 : index + %multiple_token_tiles = index.cmp ugt, %active_token_tiles, %c1 : index + %larger_shape = scalar.andi %contracts, %full_larger_channels : i1 + %larger_grid_tile = scalar.andi %larger_shape, %enough_workgroups : i1 + %larger_tile = scalar.andi %larger_grid_tile, %multiple_token_tiles : i1 + %staged_tile_channels = scf.select %larger_tile, %c256, %ordinary_tile_channels : index + %staged_workgroup_size = scf.select %larger_tile, %c1024, %ordinary_workgroup_size : index + %staged_output_tiles = index.div %output_size, %staged_tile_channels : index + %staged_token_tiles = index.div %token_capacity, %c256 : index + %input_layout = config.get @ggml.mul_mat.f16_input_layout : index + %packed = index.cmp eq, %input_layout, %c1 : index + %workgroup_size = scf.select %packed, %c512, %staged_workgroup_size : index + %token_tiles = scf.select %packed, %c1, %staged_token_tiles : index + %packed_output_tiles = index.div %output_size, %c64 : index + %output_tiles = scf.select %packed, %packed_output_tiles, %staged_output_tiles : index + kernel.launch.config workgroups(%token_tiles, %output_tiles, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 256, 2048), mul(%token_count, 256), le(%token_count, %token_capacity)] : index + %bounded_output_size = index.assume %output_size [range(%output_size, 64, 262144), mul(%output_size, 64)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %paired = scalar.constant false : i1 + %input_layout = config.get @ggml.mul_mat.f16_input_layout : index + %packed_layout = index.constant 1 : index + %packed_input = index.cmp eq, %input_layout, %packed_layout : index + %prefill_weight_format = index.constant 44 : index + %prefill_binary_op = index.constant 0 : index + func.call @ggml_mul_mat_quantized_f16_prefill_wave32(%prefill_weight_format, %prefill_binary_op, %paired, %packed_input, %paired, %bounded_token_count, %input_size, %bounded_output_size, %channel_tile, %token_tile, %input, %weight, %weight, %output, %output) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +func.decl @ggml_mul_mat_quantized_f16_prefill_conv4(%weight_format: index, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %input: buffer, %weight: buffer, %state: buffer, %filter: buffer, %output: buffer, %cache: buffer) + +kernel.def target(@ggml_mul_mat_gfx11_wave32) @ggml_mul_mat_quantized_f16_wmma_prefill_conv4(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %one = index.constant 1 : index + %sixtyfour = index.constant 64 : index + %threads = index.constant 512 : index + %groups = index.div %output_size, %sixtyfour : index + kernel.launch.config workgroups(%one, %groups, %one) workgroup_size(%threads, %one, %one) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %state: buffer, %filter: buffer, %output: buffer, %cache: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %channel_tile = kernel.workgroup.id : index + func.call @ggml_mul_mat_quantized_f16_prefill_conv4(%weight_format, %token_count, %input_size, %output_size, %channel_tile, %input, %weight, %state, %filter, %output, %cache) : (index, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +// Recurrent state is consumed by the later finish dispatch. +kernel.def target(@ggml_mul_mat_gfx11_wave32) @ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior(%token_count: index) { + %output_size = config.get @ggml.mul_mat.output_size : index + %one = index.constant 1 : index + %sixtyfour = index.constant 64 : index + %threads = index.constant 512 : index + %groups = index.div %output_size, %sixtyfour : index + kernel.launch.config workgroups(%one, %groups, %one) workgroup_size(%threads, %one, %one) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %filter: buffer, %output: buffer, %edges: buffer) { + %input_size = config.get @ggml.mul_mat.input_size : index + %output_size = config.get @ggml.mul_mat.output_size : index + %weight_format = config.get @ggml.mul_mat.weight_format : index + %channel_tile = kernel.workgroup.id : index + func.call @ggml_mul_mat_quantized_f16_prefill_conv4_interior(%weight_format, %token_count, %input_size, %output_size, %channel_tile, %input, %weight, %filter, %output, %edges) : (index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom new file mode 100644 index 000000000000..37d864977a24 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q5_k_q8_plane_wmma.loom @@ -0,0 +1,1408 @@ +template.decl @ggml.mul_mat_q8_1_x4.entry(%token_count: index, %paired_weights: i1, %q8_input: buffer, %weight: buffer, %peer_weight: buffer, %output: buffer) + +// Generic Q5_K projection consuming the Q8_1 plane published by the fused +// gate/up path. Weights use the backend-resident symmetric-I5/K32 carrier; +// input, output, and token dimensions are all JIT configuration values. + +template.decl @ggml.mul_mat_q5_k_q8_plane.wmmai8.body(%is_iq4xs: i1, %token256: i1, %q8_plane: i1, %paired_weights: i1, %peer_weight: buffer, %src0_na: buffer, %dst_na: buffer, %q8_na: buffer, %base: offset, %lds_bytes: offset, %w_off: offset, %as_off: offset, %ws_off: offset, %asum_off: offset, %wc_off: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index, %is_symi5: i1, %is_q4: i1, %is_q4_packed: i1) + +amdgpu.target @ggml_q5_plane_dense_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.mul_mat_q5_k_q8_plane.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_q5_k_q8_plane.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_q5_k_q8_plane.token_capacity : %value: index where [range(%value, 1, 2048), range(%value, 1, 2048)] + +config.decl @ggml.mul_mat_q8_1_x4.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_q8_1_x4.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.mul_mat_q8_1_x4.token_capacity : %value: index where [range(%value, 256, 2048), mul(%value, 256)] + +config.decl @ggml.mul_mat_q8_1_x4.weight_format : %value: index where [range(%value, 4, 44)] + +// Paired projections share a tile without persisting gate/up intermediates. +// Select loaded words rather than opaque buffer identities in device code. +func.def inline @ggml_q8_1_x4_paired_weight_words(%paired: i1, %load_up: i1, %weight: buffer, %peer: buffer, %offset: offset) -> (vector<4xi32>) { + %zero = index.constant 0 : index + %weight_view = buffer.view %weight[%offset] : buffer -> view<4xi32> + %weight_words = vector.load %weight_view[%zero] : view<4xi32> -> vector<4xi32> + %peer_words = scf.if %paired -> (vector<4xi32>) { + %peer_view = buffer.view %peer[%offset] : buffer -> view<4xi32> + %words = vector.load %peer_view[%zero] : view<4xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } else { + scf.yield %weight_words : vector<4xi32> + } + %words = scf.select %load_up, %peer_words, %weight_words : vector<4xi32> + func.return %words : vector<4xi32> +} + +// Decodes one Q4_K scale/minimum pair from a header loaded as three packed +// scale words. Keeping the complete header in registers lets paired-group +// contractions share both the header and packed-code load. +func.def inline @ggml_q5_plane_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c48_i32 = scalar.constant 48 : i32 + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %is_low_group = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift = index.cast %scale_shift_index : index to i32 + %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 + %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 + %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 + %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 + %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 + %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 + %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 + %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 + %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 + %scale = scalar.ori %scale_low, %scale_high : i32 + %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 + %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 + %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 + %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 + %minimum = scalar.ori %minimum_low, %minimum_high : i32 + func.return %scale, %minimum : i32, i32 +} + +// Packs either nibble from four adjacent GGUF IQ4_XS bytes for the gfx11 table lookup. +// Preserve GGUF code indices for the logical table lookup. +func.def inline @ggml_q5_plane_dense_iq4xs_wmma_table_codes4(%code_word: vector<1xi32>, %uses_high: i1) -> (vector<4xi8>) { + %c4_i32 = scalar.constant 4 : i32 + %c8_i32 = scalar.constant 8 : i32 + %c12_i32 = scalar.constant 12 : i32 + %c16_i32 = scalar.constant 16 : 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 + %word = vector.extract %code_word[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> + func.return %codes : vector<4xi8> +} + +func.def inline @ggml_q5_plane_dense_iq4xs_wmmai8_signed_values16(%code_words: vector<4xi32>, %uses_high: i1) -> (vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) { + %v0 = scalar.constant -127 : i8 + %v1 = scalar.constant -104 : i8 + %v2 = scalar.constant -83 : i8 + %v3 = scalar.constant -65 : i8 + %v4 = scalar.constant -49 : i8 + %v5 = scalar.constant -35 : i8 + %v6 = scalar.constant -22 : i8 + %v7 = scalar.constant -10 : i8 + %v8 = scalar.constant 1 : i8 + %v9 = scalar.constant 13 : i8 + %v10 = scalar.constant 25 : i8 + %v11 = scalar.constant 38 : i8 + %v12 = scalar.constant 53 : i8 + %v13 = scalar.constant 69 : i8 + %v14 = scalar.constant 89 : i8 + %v15 = scalar.constant 113 : i8 + %value_table = vector.from_elements %v0, %v1, %v2, %v3, %v4, %v5, %v6, %v7, %v8, %v9, %v10, %v11, %v12, %v13, %v14, %v15 : vector<16xi8> + %word0_i32 = vector.extract %code_words[0] : vector<4xi32> -> i32 + %word1_i32 = vector.extract %code_words[1] : vector<4xi32> -> i32 + %word2_i32 = vector.extract %code_words[2] : vector<4xi32> -> i32 + %word3_i32 = vector.extract %code_words[3] : vector<4xi32> -> i32 + %word0 = vector.from_elements %word0_i32 : vector<1xi32> + %word1 = vector.from_elements %word1_i32 : vector<1xi32> + %word2 = vector.from_elements %word2_i32 : vector<1xi32> + %word3 = vector.from_elements %word3_i32 : vector<1xi32> + %codes0 = func.call @ggml_q5_plane_dense_iq4xs_wmma_table_codes4(%word0, %uses_high) : (vector<1xi32>, i1) -> (vector<4xi8>) + %codes1 = func.call @ggml_q5_plane_dense_iq4xs_wmma_table_codes4(%word1, %uses_high) : (vector<1xi32>, i1) -> (vector<4xi8>) + %codes2 = func.call @ggml_q5_plane_dense_iq4xs_wmma_table_codes4(%word2, %uses_high) : (vector<1xi32>, i1) -> (vector<4xi8>) + %codes3 = func.call @ggml_q5_plane_dense_iq4xs_wmma_table_codes4(%word3, %uses_high) : (vector<1xi32>, i1) -> (vector<4xi8>) + %values0_i8 = vector.table.lookup %value_table[%codes0] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + %values1_i8 = vector.table.lookup %value_table[%codes1] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + %values2_i8 = vector.table.lookup %value_table[%codes2] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + %values3_i8 = vector.table.lookup %value_table[%codes3] : vector<16xi8>, vector<4xi8> -> vector<4xi8> + func.return %values0_i8, %values1_i8, %values2_i8, %values3_i8 : vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8> +} + +func.def inline @ggml_q5_plane_dense_iq4xs_wmma_scale(%weight: buffer, %row_byte_base: offset, %iq4_block: index, %iq4_group: index) -> (f32) { + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %block_bytes = index.constant 136 : offset + %scale_high_offset = index.constant 2 : offset + %scale_low_offset = index.constant 4 : offset + %c2_i32 = scalar.constant 2 : i32 + %c3_i32 = scalar.constant 3 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15_i32 = scalar.constant 15 : i32 + %c32_i32 = scalar.constant 32 : i32 + %bounded_group = index.assume %iq4_group [range(%iq4_group, 0, 7)] : index + %block_byte_add = index.scale %iq4_block, %block_bytes : index, offset -> offset + %block_byte_base = index.add %row_byte_base, %block_byte_add : offset + %scale_high_byte_base = index.add %block_byte_base, %scale_high_offset : offset + %scale_low_byte_base = index.add %block_byte_base, %scale_low_offset : offset + %d_view = buffer.view %weight[%block_byte_base] : buffer -> view<1xf16> + %scale_high_view = buffer.view %weight[%scale_high_byte_base] : buffer -> view<1xi16> + %scale_low_view = buffer.view %weight[%scale_low_byte_base] : buffer -> view<4xi8> + %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %scale_low_index = index.div %bounded_group, %c2 : index + %scale_low_i8 = view.load %scale_low_view[%scale_low_index] : view<4xi8> -> i8 + %scale_low_i32 = scalar.extui %scale_low_i8 : i8 to i32 + %scale_low_half = index.rem %bounded_group, %c2 : index + %scale_low_shift_index = index.mul %scale_low_half, %c4 : index + %scale_low_shift = index.cast %scale_low_shift_index : index to i32 + %scale_low_shifted = scalar.shrui %scale_low_i32, %scale_low_shift : i32 + %scale_low = scalar.andi %scale_low_shifted, %c15_i32 : i32 + %scale_high_i16 = view.load %scale_high_view[%c0] : view<1xi16> -> i16 + %scale_high_i32 = scalar.extui %scale_high_i16 : i16 to i32 + %scale_high_shift_index = index.mul %bounded_group, %c2 : index + %scale_high_shift = index.cast %scale_high_shift_index : index to i32 + %scale_high_shifted = scalar.shrui %scale_high_i32, %scale_high_shift : i32 + %scale_high_two_bits = scalar.andi %scale_high_shifted, %c3_i32 : i32 + %scale_high = scalar.shli %scale_high_two_bits, %c4_i32 : i32 + %scale_unsigned = scalar.ori %scale_low, %scale_high : i32 + %scale_signed = scalar.subi %scale_unsigned, %c32_i32 : i32 + %scale_f32 = scalar.sitofp %scale_signed : i32 to f32 + %d_scale = scalar.mulf %d, %scale_f32 : f32 + func.return %d_scale : f32 +} + +// Token-128 and token-256 share one body; a compile-time selector removes the unused statically bounded LDS view. +func.def inline @ggml_q5_plane_dense_q8_wmmai8_stage_values(%token256: i1, %scratch: buffer, %base: offset, %adbase: index, %ad1: index, %values0: vector<32xi8>, %values1: vector<32xi8>) { + scf.if %token256 { + %al_flat = buffer.view %scratch[%base] : buffer -> view<20480xi8> + %ad0_bounded = index.assume %adbase [range(%adbase, 0, 20400)] : index + %ad1_bounded = index.assume %ad1 [range(%ad1, 32, 20432)] : index + vector.store %values0, %al_flat[%ad0_bounded] : vector<32xi8>, view<20480xi8> + vector.store %values1, %al_flat[%ad1_bounded] : vector<32xi8>, view<20480xi8> + } else { + %al_flat = buffer.view %scratch[%base] : buffer -> view<10240xi8> + %ad0_bounded = index.assume %adbase [range(%adbase, 0, 10160)] : index + %ad1_bounded = index.assume %ad1 [range(%ad1, 32, 10192)] : index + vector.store %values0, %al_flat[%ad0_bounded] : vector<32xi8>, view<10240xi8> + vector.store %values1, %al_flat[%ad1_bounded] : vector<32xi8>, view<10240xi8> + } + func.return +} + +func.def inline @ggml_q5_plane_dense_q8_wmmai8_stage_metadata(%token256: i1, %scratch: buffer, %as_off: offset, %asum_off: offset, %dst0: index, %dst1: index, %d0: f32, %d1: f32, %sum0: f32, %sum1: f32) { + scf.if %token256 { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<512xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<512xf32> + %dst0_bounded = index.assume %dst0 [range(%dst0, 0, 495)] : index + %dst1_bounded = index.assume %dst1 [range(%dst1, 16, 511)] : index + view.store %d0, %scale_view[%dst0_bounded] : f32, view<512xf32> + view.store %d1, %scale_view[%dst1_bounded] : f32, view<512xf32> + view.store %sum0, %sum_view[%dst0_bounded] : f32, view<512xf32> + view.store %sum1, %sum_view[%dst1_bounded] : f32, view<512xf32> + } else { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<256xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<256xf32> + %dst0_bounded = index.assume %dst0 [range(%dst0, 0, 239)] : index + %dst1_bounded = index.assume %dst1 [range(%dst1, 16, 255)] : index + view.store %d0, %scale_view[%dst0_bounded] : f32, view<256xf32> + view.store %d1, %scale_view[%dst1_bounded] : f32, view<256xf32> + view.store %sum0, %sum_view[%dst0_bounded] : f32, view<256xf32> + view.store %sum1, %sum_view[%dst1_bounded] : f32, view<256xf32> + } + func.return +} + +func.def inline @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256: i1, %scratch: buffer, %base: offset, %row: index, %column: index, %c16: index) -> (vector<16xi8>) { + %lhs_layout = encoding.layout.strided [80, 1] : encoding + %result = scf.if %token256 -> (vector<16xi8>) { + %al_view = buffer.view %scratch[%base] : buffer -> view<256x64xi8, %lhs_layout> + %values = vector.fragment.load %al_view[%row, %column] shape [%c16, %c16] : view<256x64xi8, %lhs_layout> -> vector<16xi8> + scf.yield %values : vector<16xi8> + } else { + %al_view = buffer.view %scratch[%base] : buffer -> view<128x64xi8, %lhs_layout> + %values = vector.fragment.load %al_view[%row, %column] shape [%c16, %c16] : view<128x64xi8, %lhs_layout> -> vector<16xi8> + scf.yield %values : vector<16xi8> + } + func.return %result : vector<16xi8> +} + +func.def inline @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256: i1, %scratch: buffer, %as_off: offset, %asum_off: offset, %index: index) -> (vector<8xf32>, vector<8xf32>) { + %scale, %sum = scf.if %token256 -> (vector<8xf32>, vector<8xf32>) { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<512xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<512xf32> + %bounded_index = index.assume %index [range(%index, 0, 504)] : index + %scale_values = vector.load %scale_view[%bounded_index] : view<512xf32> -> vector<8xf32> + %sum_values = vector.load %sum_view[%bounded_index] : view<512xf32> -> vector<8xf32> + scf.yield %scale_values, %sum_values : vector<8xf32>, vector<8xf32> + } else { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<256xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<256xf32> + %bounded_index = index.assume %index [range(%index, 0, 120)] : index + %scale_values = vector.load %scale_view[%bounded_index] : view<256xf32> -> vector<8xf32> + %sum_values = vector.load %sum_view[%bounded_index] : view<256xf32> -> vector<8xf32> + scf.yield %scale_values, %sum_values : vector<8xf32>, vector<8xf32> + } + func.return %scale, %sum : vector<8xf32>, vector<8xf32> +} + +// Raw Q5_K x packed-Q8_1_x4 IU8 WMMA using the existing Loom activation layout. +// Keep one weighted K32 fragment live while updating the paired projection tile. +func.def inline @ggml_q8_1_x4_affine_rhs_tile(%scratch: buffer, %weight_offset: offset, %scale_offset: offset, %correction_offset: offset, %block: index, %column: index, %lane: index, %lhs0: vector<4xi32>, %lhs1: vector<4xi32>, %zero: vector<8xi32>, %activation_scale: vector<8xf32>, %activation_sum: vector<8xf32>, %accumulator: vector<8xf32>) -> (vector<8xf32>) { + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %row_layout = encoding.layout.strided [1, 80] : encoding + %rhs_schema = encoding.define #encoding.operand : encoding + %weight = buffer.view %scratch[%weight_offset] : buffer -> view<64x128xi8, %row_layout> + %scales = buffer.view %scratch[%scale_offset] : buffer -> view<256xf32> + %corrections = buffer.view %scratch[%correction_offset] : buffer -> view<256xf32> + %row = index.add %column, %lane : index + %row_pair = index.mul %row, %c2 : index + %metadata0 = index.add %row_pair, %block : index + %metadata = index.assume %metadata0 [range(%metadata0, 0, 255)] : index + %scale = view.load %scales[%metadata] : view<256xf32> -> f32 + %scale_vector = vector.splat %scale : vector<8xf32> + %correction = view.load %corrections[%metadata] : view<256xf32> -> f32 + %correction_vector = vector.splat %correction : vector<8xf32> + %k0 = index.mul %block, %c32 : index + %k1 = index.add %k0, %c16 : index + %rhs0_raw = vector.fragment.load %weight[%k0, %column] shape [%c16, %c16] : view<64x128xi8, %row_layout> -> vector<16xi8> + %rhs0_words = vector.bitcast %rhs0_raw : vector<16xi8> to vector<4xi32> + %rhs0 = vector.fragment %rhs0_words shape [%c16, %c16] using {schema = %rhs_schema : encoding} : vector<4xi32> + %rhs1_raw = vector.fragment.load %weight[%k1, %column] shape [%c16, %c16] : view<64x128xi8, %row_layout> -> vector<16xi8> + %rhs1_words = vector.bitcast %rhs1_raw : vector<16xi8> to vector<4xi32> + %rhs1 = vector.fragment %rhs1_words shape [%c16, %c16] using {schema = %rhs_schema : encoding} : vector<4xi32> + %dot0 = vector.mma %lhs0, %rhs0, %zero : vector<4xi32>, vector<4xi32>, vector<8xi32> + %dot = vector.mma %lhs1, %rhs1, %dot0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %dot_f32 = vector.sitofp %dot : vector<8xi32> to vector<8xf32> + %product_scale = vector.mulf %activation_scale, %scale_vector : vector<8xf32> + %base = vector.fmaf %activation_sum, %correction_vector, %accumulator : vector<8xf32> + %result = vector.fmaf %dot_f32, %product_scale, %base : vector<8xf32> + func.return %result : vector<8xf32> +} + +template.def<@ggml.mul_mat_q5_k_q8_plane.wmmai8.body> device @ggml_mul_mat_q5_k_q8_plane_wmmai8_body(%is_iq4xs: i1, %token256: i1, %q8_plane: i1, %paired_weights: i1, %peer_weight: buffer, %src0_na: buffer, %dst_na: buffer, %q8_na: buffer, %base: offset, %lds_bytes: offset, %w_off: offset, %as_off: offset, %ws_off: offset, %asum_off: offset, %wc_off: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index, %is_symi5: i1, %is_q4: i1, %is_q4_packed: i1) { + %c64_wide = index.constant 64 : index + %weight_elements = index.mul %k_n, %k_wstride : index + %weight_metadata = index.mul %k_n, %c2 : index + %q4_k_block = index.constant 256 : index + %q5_native_block_bytes = index.constant 176 : offset + %q4_block_bytes = index.constant 144 : offset + %q5_block_bytes = scf.select %is_q4, %q4_block_bytes, %q5_native_block_bytes : offset + %q5_high_offset = index.constant 16 : offset + %q5_native_code_offset = index.constant 48 : offset + %q4_code_offset = index.constant 16 : offset + %q5_code_offset = scf.select %is_q4, %q4_code_offset, %q5_native_code_offset : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_offset = index.constant 16 : offset + %q8_elements_per_group = index.constant 128 : index + // dst is [rows, cols] with rows contiguous, and C is [cols, rows]. + %dst_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + + // Waves share the staged weight tile across their token rows. + %scratch = buffer.alloca align(16) %lds_bytes : buffer + %i8_schema = encoding.define #encoding.operand : encoding + // Compile-time formats select unsigned Q5_K or signed IQ4_XS/I5 operands. + %signed_rhs_schema = encoding.define #encoding.operand : encoding + %unsigned_rhs_schema = encoding.define #encoding.operand : encoding + %is_signed_rhs = scalar.ori %is_iq4xs, %is_symi5 : i1 + %u8_schema = scf.select %is_signed_rhs, %signed_rhs_schema, %unsigned_rhs_schema : encoding + %rhs_layout = encoding.layout.strided [1, 80] : encoding + %wl_view = buffer.view %scratch[%w_off] : buffer -> view<64x[%k_n]xi8, %rhs_layout> + %wl_flat = buffer.view %scratch[%w_off] : buffer -> view<[%weight_elements]xi8> + %wsl_view = buffer.view %scratch[%ws_off] : buffer -> view<[%weight_metadata]xf32> + %wcl_view = buffer.view %scratch[%wc_off] : buffer -> view<[%weight_metadata]xf32> + + %col_base = index.mul %col_tile, %k_m : index + %output_tile_rows = scf.select %paired_weights, %c64_wide, %k_n : index + %row_base = index.mul %row_tile, %output_tile_rows : index + %up_half = index.cmp uge, %tid, %c64_wide : index + %load_up = scalar.andi %paired_weights, %up_half : i1 + + // One thread stages both K32 blocks for a token from the 144-byte Q8_1_x4 group. + %ascol_g = index.add %col_base, %tid : index + %adbase = index.mul %tid, %k_astride : index + %q8_groups_per_row = index.div %k_b, %q8_elements_per_group : index + %q8_row_bytes = index.scale %q8_groups_per_row, %q8_group_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %ascol_g, %q8_row_bytes : index, offset -> offset + %q8_plane_payload_elements = index.mul %k_b, %cols_b : index + %q8_plane_payload_words = index.div %q8_plane_payload_elements, %c4 : index + %q8_plane_metadata_offset = index.cast %q8_plane_payload_elements : index to offset + %q8_plane_metadata_count = index.mul %kblocks, %cols_b : index + %q8_plane_payload = buffer.view %q8_na[%base] : buffer -> view<[%q8_plane_payload_words]xi32> + %q8_plane_metadata = buffer.view %q8_na[%q8_plane_metadata_offset] : buffer -> view<[%q8_plane_metadata_count]xi32> + + // Paired projections use fewer token rows and a wider shared weight tile. + %wcol = index.rem %wave, %k_wm : index + %wrow = index.div %wave, %k_wm : index + %k_mspan = scf.select %paired_weights, %c16, %c32 : index + %k_nspan = index.add %k_n, %c0 : index + %wm_off = index.mul %wcol, %k_mspan : index + %wn_off = index.mul %wrow, %k_nspan : index + %m_out = index.add %col_base, %wm_off : index + %n_out = index.add %row_base, %wn_off : index + %lm0 = index.add %wm_off, %c0 : index + %gm0 = index.add %m_out, %c0 : index + %k_ma1 = index.constant 16 : index + %lm1 = index.add %wm_off, %k_ma1 : index + %gm1 = index.add %m_out, %k_ma1 : index + %ln0 = index.add %wn_off, %c0 : index + %gn0 = index.add %n_out, %c0 : index + %k_nb1 = index.constant 16 : index + %ln1 = index.add %wn_off, %k_nb1 : index + %gn1 = index.add %n_out, %k_nb1 : index + %k_nb2 = index.constant 32 : index + %ln2 = index.add %wn_off, %k_nb2 : index + %gn2 = index.add %n_out, %k_nb2 : index + %k_nb3 = index.constant 48 : index + %ln3 = index.add %wn_off, %k_nb3 : index + %gn3 = index.add %n_out, %k_nb3 : index + %lane = index.rem %tid, %k_wave : index + %lane_lo = index.rem %lane, %c16 : index + %lane_hi = index.div %lane, %c16 : index + %wsb0_0 = index.add %ln0, %lane_lo : index + %wsb0 = index.mul %wsb0_0, %k_blocks : index + %wsb1_0 = index.add %ln1, %lane_lo : index + %wsb1 = index.mul %wsb1_0, %k_blocks : index + %wsb2_0 = index.add %ln2, %lane_lo : index + %wsb2 = index.mul %wsb2_0, %k_blocks : index + %wsb3_0 = index.add %ln3, %lane_lo : index + %wsb3 = index.mul %wsb3_0, %k_blocks : index + // Activation scale metadata is [tile][block][lane-half][register]. + %as_tile0 = index.div %lm0, %c16 : index + %asb0 = index.mul %as_tile0, %k_blocks : index + %as_tile1 = index.div %lm1, %c16 : index + %asb1 = index.mul %as_tile1, %k_blocks : index + + %f0_0, %f0_1, %f0_2, %f0_3, %f1_0, %f1_1, %f1_2, %f1_3 = scf.for %chunk = [%c0 to %nchunks step %c1](%fc0_0 = %fzero : vector<8xf32>, %fc0_1 = %fzero : vector<8xf32>, %fc0_2 = %fzero : vector<8xf32>, %fc0_3 = %fzero : vector<8xf32>, %fc1_0 = %fzero : vector<8xf32>, %fc1_1 = %fzero : vector<8xf32>, %fc1_2 = %fzero : vector<8xf32>, %fc1_3 = %fzero : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %activation_active = index.cmp ult, %tid, %k_m : index + scf.if %activation_active { + %q8_words0, %q8_words1, %q8_d0_f16, %q8_s0_f16, %q8_d1_f16, %q8_s1_f16 = scf.if %q8_plane -> (vector<8xi32>, vector<8xi32>, f16, f16, f16, f16) { + %q8_block0 = index.mul %chunk, %c2 : index + %q8_block1 = index.add %q8_block0, %c1 : index + %q8_block0_token_base = index.mul %q8_block0, %cols_b : index + %q8_block1_token_base = index.mul %q8_block1, %cols_b : index + %q8_block0_token0 = index.add %q8_block0_token_base, %ascol_g : index + %q8_block1_token0 = index.add %q8_block1_token_base, %ascol_g : index + %q8_block0_token = index.assume %q8_block0_token0 [range(%q8_block0_token0, 0, 33554431)] : index + %q8_block1_token = index.assume %q8_block1_token0 [range(%q8_block1_token0, 1, 33554431)] : index + %q8_payload_word0_0 = index.mul %q8_block0_token, %c8 : index + %q8_payload_word1_0 = index.mul %q8_block1_token, %c8 : index + %q8_payload_word0 = index.assume %q8_payload_word0_0 [range(%q8_payload_word0_0, 0, 268435448), mul(%q8_payload_word0_0, 8)] : index + %q8_payload_word1 = index.assume %q8_payload_word1_0 [range(%q8_payload_word1_0, 8, 268435448), mul(%q8_payload_word1_0, 8)] : index + %q8_plane_words0 = vector.load %q8_plane_payload[%q8_payload_word0] : view<[%q8_plane_payload_words]xi32> -> vector<8xi32> + %q8_plane_words1 = vector.load %q8_plane_payload[%q8_payload_word1] : view<[%q8_plane_payload_words]xi32> -> vector<8xi32> + %q8_meta0_word = view.load %q8_plane_metadata[%q8_block0_token] : view<[%q8_plane_metadata_count]xi32> -> i32 + %q8_meta1_word = view.load %q8_plane_metadata[%q8_block1_token] : view<[%q8_plane_metadata_count]xi32> -> i32 + %q8_meta0_vector = vector.from_elements %q8_meta0_word : vector<1xi32> + %q8_meta1_vector = vector.from_elements %q8_meta1_word : vector<1xi32> + %q8_ds0_vector = vector.bitcast %q8_meta0_vector : vector<1xi32> to vector<2xf16> + %q8_ds1_vector = vector.bitcast %q8_meta1_vector : vector<1xi32> to vector<2xf16> + %q8_plane_d0_f16 = vector.extract %q8_ds0_vector[0] : vector<2xf16> -> f16 + %q8_plane_s0_f16 = vector.extract %q8_ds0_vector[1] : vector<2xf16> -> f16 + %q8_plane_d1_f16 = vector.extract %q8_ds1_vector[0] : vector<2xf16> -> f16 + %q8_plane_s1_f16 = vector.extract %q8_ds1_vector[1] : vector<2xf16> -> f16 + scf.yield %q8_plane_words0, %q8_plane_words1, %q8_plane_d0_f16, %q8_plane_s0_f16, %q8_plane_d1_f16, %q8_plane_s1_f16 : vector<8xi32>, vector<8xi32>, f16, f16, f16, f16 + } else { + %q8_group = index.div %chunk, %c2 : index + %q8_pair = index.rem %chunk, %c2 : index + %q8_inner0 = index.mul %q8_pair, %c2 : index + %q8_inner1 = index.add %q8_inner0, %c1 : index + %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset + %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset + %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_offset : offset + %q8_ds = buffer.view %q8_na[%q8_group_byte_base] : buffer -> view<8xf16> + %q8_payload = buffer.view %q8_na[%q8_payload_byte_base] : buffer -> view<32xi32> + %q8_word0 = index.mul %q8_inner0, %c8 : index + %q8_word1 = index.mul %q8_inner1, %c8 : index + %q8_standard_words0 = vector.load %q8_payload[%q8_word0] : view<32xi32> -> vector<8xi32> + %q8_standard_words1 = vector.load %q8_payload[%q8_word1] : view<32xi32> -> vector<8xi32> + %q8_ds0 = index.mul %q8_inner0, %c2 : index + %q8_ds1 = index.mul %q8_inner1, %c2 : index + %q8_s0 = index.add %q8_ds0, %c1 : index + %q8_s1 = index.add %q8_ds1, %c1 : index + %q8_standard_d0_f16 = view.load %q8_ds[%q8_ds0] : view<8xf16> -> f16 + %q8_standard_s0_f16 = view.load %q8_ds[%q8_s0] : view<8xf16> -> f16 + %q8_standard_d1_f16 = view.load %q8_ds[%q8_ds1] : view<8xf16> -> f16 + %q8_standard_s1_f16 = view.load %q8_ds[%q8_s1] : view<8xf16> -> f16 + scf.yield %q8_standard_words0, %q8_standard_words1, %q8_standard_d0_f16, %q8_standard_s0_f16, %q8_standard_d1_f16, %q8_standard_s1_f16 : vector<8xi32>, vector<8xi32>, f16, f16, f16, f16 + } + %q8_values0 = vector.bitcast %q8_words0 : vector<8xi32> to vector<32xi8> + %q8_values1 = vector.bitcast %q8_words1 : vector<8xi32> to vector<32xi8> + %ad1_0 = index.add %adbase, %c32 : index + func.call @ggml_q5_plane_dense_q8_wmmai8_stage_values(%token256, %scratch, %base, %adbase, %ad1_0, %q8_values0, %q8_values1) : (i1, buffer, offset, index, index, vector<32xi8>, vector<32xi8>) + + // Match the scale-plane permutation consumed by the WMMA fragments. + %q8_token_tile = index.div %tid, %c16 : index + %q8_token_inner = index.rem %tid, %c16 : index + %q8_token_parity = index.rem %q8_token_inner, %c2 : index + %q8_token_vector = index.div %q8_token_inner, %c2 : index + %q8_scale_tile0 = index.mul %q8_token_tile, %c4 : index + %q8_scale_tile1 = index.add %q8_scale_tile0, %q8_token_parity : index + %q8_scale_tile2 = index.mul %q8_scale_tile1, %c8 : index + %q8_scale_dst0_0 = index.add %q8_scale_tile2, %q8_token_vector : index + %q8_scale_dst1_0 = index.add %q8_scale_dst0_0, %c16 : index + %q8_d0 = scalar.extf %q8_d0_f16 : f16 to f32 + %q8_scaled_sum0 = scalar.extf %q8_s0_f16 : f16 to f32 + %q8_d1 = scalar.extf %q8_d1_f16 : f16 to f32 + %q8_scaled_sum1 = scalar.extf %q8_s1_f16 : f16 to f32 + func.call @ggml_q5_plane_dense_q8_wmmai8_stage_metadata(%token256, %scratch, %as_off, %asum_off, %q8_scale_dst0_0, %q8_scale_dst1_0, %q8_d0, %q8_d1, %q8_scaled_sum0, %q8_scaled_sum1) : (i1, buffer, offset, offset, index, index, f32, f32, f32, f32) + } + // One lane decodes two adjacent Q5_K groups; each selects its bits from the shared 32-byte fifth-bit plane. + %q5_weight_lane_count = index.add %k_n, %c0 : index + %q5_weight_lane_active = index.cmp ult, %tid, %q5_weight_lane_count : index + scf.if %q5_weight_lane_active { + scf.if %is_iq4xs { + %iq4_block_bytes = index.constant 136 : offset + %iq4_code_offset = index.constant 8 : offset + %iq4_weight_tid0 = index.rem %tid, %q5_weight_lane_count : index + %iq4_weight_tid = index.assume %iq4_weight_tid0 [range(%iq4_weight_tid0, 0, 63)] : index + %iq4_block_count = index.div %k_b, %q4_k_block : index + %iq4_block = index.div %chunk, %c4 : index + %iq4_pair = index.rem %chunk, %c4 : index + %iq4_group0 = index.mul %iq4_pair, %c2 : index + %iq4_group1 = index.add %iq4_group0, %c1 : index + %iq4_row0 = index.add %row_base, %iq4_weight_tid : index + %iq4_row = index.assume %iq4_row0 [range(%iq4_row0, 0, 262143)] : index + %iq4_row_bytes = index.scale %iq4_block_count, %iq4_block_bytes : index, offset -> offset + %iq4_row_byte_base = index.scale %iq4_row, %iq4_row_bytes : index, offset -> offset + %iq4_block_byte_add = index.scale %iq4_block, %iq4_block_bytes : index, offset -> offset + %iq4_block_byte_base = index.add %iq4_row_byte_base, %iq4_block_byte_add : offset + %iq4_code_byte_base = index.add %iq4_block_byte_base, %iq4_code_offset : offset + %iq4_code_view = buffer.view %src0_na[%iq4_code_byte_base] : buffer -> view<32xi32> + %iq4_scale0 = func.call @ggml_q5_plane_dense_iq4xs_wmma_scale(%src0_na, %iq4_row_byte_base, %iq4_block, %iq4_group0) : (buffer, offset, index, index) -> (f32) + %iq4_scale1 = func.call @ggml_q5_plane_dense_iq4xs_wmma_scale(%src0_na, %iq4_row_byte_base, %iq4_block, %iq4_group1) : (buffer, offset, index, index) -> (f32) + %iq4_meta0_0 = index.mul %iq4_weight_tid, %c2 : index + %iq4_meta0 = index.assume %iq4_meta0_0 [range(%iq4_meta0_0, 0, 126)] : index + %iq4_meta1_0 = index.add %iq4_meta0, %c1 : index + %iq4_meta1 = index.assume %iq4_meta1_0 [range(%iq4_meta1_0, 1, 127)] : index + %iq4_bias = scalar.constant -128.0 : f32 + %iq4_correction0 = scalar.mulf %iq4_scale0, %iq4_bias : f32 + %iq4_correction1 = scalar.mulf %iq4_scale1, %iq4_bias : f32 + view.store %iq4_scale0, %wsl_view[%iq4_meta0] : f32, view<[%weight_metadata]xf32> + view.store %iq4_scale1, %wsl_view[%iq4_meta1] : f32, view<[%weight_metadata]xf32> + view.store %iq4_correction0, %wcl_view[%iq4_meta0] : f32, view<[%weight_metadata]xf32> + view.store %iq4_correction1, %wcl_view[%iq4_meta1] : f32, view<[%weight_metadata]xf32> + + %iq4_group0_word_base0 = index.mul %iq4_group0, %c4 : index + %iq4_group0_word_base = index.assume %iq4_group0_word_base0 [range(%iq4_group0_word_base0, 0, 28)] : index + %iq4_group1_word_base0 = index.mul %iq4_group1, %c4 : index + %iq4_group1_word_base = index.assume %iq4_group1_word_base0 [range(%iq4_group1_word_base0, 4, 28)] : index + %iq4_weight_row_base = index.mul %iq4_weight_tid, %k_wstride : index + %iq4_w0 = index.assume %iq4_weight_row_base [range(%iq4_weight_row_base, 0, 5040)] : index + %iq4_payload16 = index.constant 16 : index + %iq4_payload48 = index.constant 48 : index + %iq4_w0_hi0 = index.add %iq4_weight_row_base, %iq4_payload16 : index + %iq4_w0_hi = index.assume %iq4_w0_hi0 [range(%iq4_w0_hi0, 16, 5056)] : index + %iq4_w1_0 = index.add %iq4_weight_row_base, %c32 : index + %iq4_w1 = index.assume %iq4_w1_0 [range(%iq4_w1_0, 32, 5072)] : index + %iq4_w1_hi0 = index.add %iq4_weight_row_base, %iq4_payload48 : index + %iq4_w1_hi = index.assume %iq4_w1_hi0 [range(%iq4_w1_hi0, 48, 5088)] : index + %iq4_false = scalar.constant false : i1 + %iq4_true = scalar.constant true : i1 + + %iq4_words0 = vector.load %iq4_code_view[%iq4_group0_word_base] : view<32xi32> -> vector<4xi32> + %iq4_values00, %iq4_values01, %iq4_values02, %iq4_values03 = func.call @ggml_q5_plane_dense_iq4xs_wmmai8_signed_values16(%iq4_words0, %iq4_false) : (vector<4xi32>, i1) -> (vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) + vector.store %iq4_values00, %wl_flat[%iq4_w0] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w0_4_0 = index.add %iq4_w0, %c4 : index + %iq4_w0_4 = index.assume %iq4_w0_4_0 [range(%iq4_w0_4_0, 4, 5044)] : index + vector.store %iq4_values01, %wl_flat[%iq4_w0_4] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w0_8_0 = index.add %iq4_w0, %c8 : index + %iq4_w0_8 = index.assume %iq4_w0_8_0 [range(%iq4_w0_8_0, 8, 5048)] : index + vector.store %iq4_values02, %wl_flat[%iq4_w0_8] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w0_12_0 = index.add %iq4_w0_8, %c4 : index + %iq4_w0_12 = index.assume %iq4_w0_12_0 [range(%iq4_w0_12_0, 12, 5052)] : index + vector.store %iq4_values03, %wl_flat[%iq4_w0_12] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_values04, %iq4_values05, %iq4_values06, %iq4_values07 = func.call @ggml_q5_plane_dense_iq4xs_wmmai8_signed_values16(%iq4_words0, %iq4_true) : (vector<4xi32>, i1) -> (vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) + vector.store %iq4_values04, %wl_flat[%iq4_w0_hi] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w0_hi4_0 = index.add %iq4_w0_hi, %c4 : index + %iq4_w0_hi4 = index.assume %iq4_w0_hi4_0 [range(%iq4_w0_hi4_0, 20, 5060)] : index + vector.store %iq4_values05, %wl_flat[%iq4_w0_hi4] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w0_hi8_0 = index.add %iq4_w0_hi, %c8 : index + %iq4_w0_hi8 = index.assume %iq4_w0_hi8_0 [range(%iq4_w0_hi8_0, 24, 5064)] : index + vector.store %iq4_values06, %wl_flat[%iq4_w0_hi8] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w0_hi12_0 = index.add %iq4_w0_hi8, %c4 : index + %iq4_w0_hi12 = index.assume %iq4_w0_hi12_0 [range(%iq4_w0_hi12_0, 28, 5068)] : index + vector.store %iq4_values07, %wl_flat[%iq4_w0_hi12] : vector<4xi8>, view<[%weight_elements]xi8> + + %iq4_words1 = vector.load %iq4_code_view[%iq4_group1_word_base] : view<32xi32> -> vector<4xi32> + %iq4_values10, %iq4_values11, %iq4_values12, %iq4_values13 = func.call @ggml_q5_plane_dense_iq4xs_wmmai8_signed_values16(%iq4_words1, %iq4_false) : (vector<4xi32>, i1) -> (vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) + vector.store %iq4_values10, %wl_flat[%iq4_w1] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w1_4_0 = index.add %iq4_w1, %c4 : index + %iq4_w1_4 = index.assume %iq4_w1_4_0 [range(%iq4_w1_4_0, 36, 5076)] : index + vector.store %iq4_values11, %wl_flat[%iq4_w1_4] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w1_8_0 = index.add %iq4_w1, %c8 : index + %iq4_w1_8 = index.assume %iq4_w1_8_0 [range(%iq4_w1_8_0, 40, 5080)] : index + vector.store %iq4_values12, %wl_flat[%iq4_w1_8] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w1_12_0 = index.add %iq4_w1_8, %c4 : index + %iq4_w1_12 = index.assume %iq4_w1_12_0 [range(%iq4_w1_12_0, 44, 5084)] : index + vector.store %iq4_values13, %wl_flat[%iq4_w1_12] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_values14, %iq4_values15, %iq4_values16, %iq4_values17 = func.call @ggml_q5_plane_dense_iq4xs_wmmai8_signed_values16(%iq4_words1, %iq4_true) : (vector<4xi32>, i1) -> (vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) + vector.store %iq4_values14, %wl_flat[%iq4_w1_hi] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w1_hi4_0 = index.add %iq4_w1_hi, %c4 : index + %iq4_w1_hi4 = index.assume %iq4_w1_hi4_0 [range(%iq4_w1_hi4_0, 52, 5092)] : index + vector.store %iq4_values15, %wl_flat[%iq4_w1_hi4] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w1_hi8_0 = index.add %iq4_w1_hi, %c8 : index + %iq4_w1_hi8 = index.assume %iq4_w1_hi8_0 [range(%iq4_w1_hi8_0, 56, 5096)] : index + vector.store %iq4_values16, %wl_flat[%iq4_w1_hi8] : vector<4xi8>, view<[%weight_elements]xi8> + %iq4_w1_hi12_0 = index.add %iq4_w1_hi8, %c4 : index + %iq4_w1_hi12 = index.assume %iq4_w1_hi12_0 [range(%iq4_w1_hi12_0, 60, 5100)] : index + vector.store %iq4_values17, %wl_flat[%iq4_w1_hi12] : vector<4xi8>, view<[%weight_elements]xi8> + } else { + %q5_weight_tid0 = index.rem %tid, %q5_weight_lane_count : index + %q5_weight_tid = index.assume %q5_weight_tid0 [range(%q5_weight_tid0, 0, 127)] : index + %q5_block_count = index.div %k_b, %q4_k_block : index + %q5_block = index.div %chunk, %c4 : index + %q5_pair = index.rem %chunk, %c4 : index + %q5_group0 = index.mul %q5_pair, %c2 : index + %q5_group1 = index.add %q5_group0, %c1 : index + %paired_row_lane = index.rem %q5_weight_tid, %c64_wide : index + %weight_row_lane = scf.select %paired_weights, %paired_row_lane, %q5_weight_tid : index + %q5_row0 = index.add %row_base, %weight_row_lane : index + %q5_row = index.assume %q5_row0 [range(%q5_row0, 0, 262143)] : index + %q5_row_bytes = index.scale %q5_block_count, %q5_block_bytes : index, offset -> offset + %q5_row_byte_base = index.scale %q5_row, %q5_row_bytes : index, offset -> offset + %q5_block_byte_add = index.scale %q5_block, %q5_block_bytes : index, offset -> offset + %q5_native_block_byte_base = index.add %q5_row_byte_base, %q5_block_byte_add : offset + %q4_row64 = index.constant 64 : index + %q4_field_bytes = index.constant 16 : offset + %q4_group_bytes = index.constant 9216 : offset + %q4_row_group = index.div %q5_row, %q4_row64 : index + %q4_record = index.madd %q4_row_group, %q5_block_count, %q5_block : index + %q4_record_base = index.scale %q4_record, %q4_group_bytes : index, offset -> offset + %packed_row_lane = index.rem %q5_row, %q4_row64 : index + %q4_row_offset = index.scale %packed_row_lane, %q4_field_bytes : index, offset -> offset + %q4_packed_block_byte_base = index.add %q4_record_base, %q4_row_offset : offset + %q5_block_byte_base = scf.select %is_q4_packed, %q4_packed_block_byte_base, %q5_native_block_byte_base : offset + %q5_high_byte_base = index.add %q5_block_byte_base, %q5_high_offset : offset + %q5_code_byte_base = index.add %q5_block_byte_base, %q5_code_offset : offset + %q5_high_view = buffer.view %src0_na[%q5_high_byte_base] : buffer -> view<8xi32> + %q5_code_view = buffer.view %src0_na[%q5_code_byte_base] : buffer -> view<32xi32> + %q5_d0, %q5_d1, %q5_correction0, %q5_correction1 = scf.if %is_symi5 -> (f32, f32, f32, f32) { + // The capacity-neutral signed-I5 carrier reuses the Q5 bit planes + // and stores one adjacent F16 scale pair for these two K32 groups. + %q5_s5_header_view = buffer.view %src0_na[%q5_block_byte_base] : buffer -> view<16xi8> + %q5_s5_scale_byte0 = index.mul %q5_group0, %c2 : index + %q5_s5_scale_byte = index.assume %q5_s5_scale_byte0 [range(%q5_s5_scale_byte0, 0, 12), mul(%q5_s5_scale_byte0, 4)] : index + %q5_s5_scale_raw = vector.load %q5_s5_header_view[%q5_s5_scale_byte] : view<16xi8> -> vector<4xi8> + %q5_s5_scale_halves = vector.bitcast %q5_s5_scale_raw : vector<4xi8> to vector<2xf16> + %q5_s5_scale0_f16 = vector.extract %q5_s5_scale_halves[0] : vector<2xf16> -> f16 + %q5_s5_scale1_f16 = vector.extract %q5_s5_scale_halves[1] : vector<2xf16> -> f16 + %q5_s5_scale0 = scalar.extf %q5_s5_scale0_f16 : f16 to f32 + %q5_s5_scale1 = scalar.extf %q5_s5_scale1_f16 : f16 to f32 + %q5_s5_zero = scalar.constant 0.0 : f32 + scf.yield %q5_s5_scale0, %q5_s5_scale1, %q5_s5_zero, %q5_s5_zero : f32, f32, f32, f32 + } else { + %q5_header = func.call @ggml_q8_1_x4_paired_weight_words(%paired_weights, %load_up, %src0_na, %peer_weight, %q5_block_byte_base) : (i1, i1, buffer, buffer, offset) -> (vector<4xi32>) + %q5_header_halves = vector.bitcast %q5_header : vector<4xi32> to vector<8xf16> + %q5_d_f16 = vector.extract %q5_header_halves[0] : vector<8xf16> -> f16 + %q5_dmin_f16 = vector.extract %q5_header_halves[1] : vector<8xf16> -> f16 + %q5_scale0 = vector.extract %q5_header[1] : vector<4xi32> -> i32 + %q5_scale1 = vector.extract %q5_header[2] : vector<4xi32> -> i32 + %q5_scale2 = vector.extract %q5_header[3] : vector<4xi32> -> i32 + %q5_d = scalar.extf %q5_d_f16 : f16 to f32 + %q5_dmin = scalar.extf %q5_dmin_f16 : f16 to f32 + %q5_scale_group0, %q5_min_group0 = func.call @ggml_q5_plane_q4k_scale_from_header(%q5_scale0, %q5_scale1, %q5_scale2, %q5_group0) : (i32, i32, i32, index) -> (i32, i32) + %q5_scale_group1, %q5_min_group1 = func.call @ggml_q5_plane_q4k_scale_from_header(%q5_scale0, %q5_scale1, %q5_scale2, %q5_group1) : (i32, i32, i32, index) -> (i32, i32) + %q5_scale_group0_f32 = scalar.uitofp %q5_scale_group0 : i32 to f32 + %q5_scale_group1_f32 = scalar.uitofp %q5_scale_group1 : i32 to f32 + %q5_min_group0_f32 = scalar.uitofp %q5_min_group0 : i32 to f32 + %q5_min_group1_f32 = scalar.uitofp %q5_min_group1 : i32 to f32 + %q5_exact_d0 = scalar.mulf %q5_d, %q5_scale_group0_f32 : f32 + %q5_exact_d1 = scalar.mulf %q5_d, %q5_scale_group1_f32 : f32 + %q5_m0 = scalar.mulf %q5_dmin, %q5_min_group0_f32 : f32 + %q5_m1 = scalar.mulf %q5_dmin, %q5_min_group1_f32 : f32 + %q5_neg_m0 = scalar.negf %q5_m0 : f32 + %q5_neg_m1 = scalar.negf %q5_m1 : f32 + scf.yield %q5_exact_d0, %q5_exact_d1, %q5_neg_m0, %q5_neg_m1 : f32, f32, f32, f32 + } + %q5_meta0_0 = index.mul %q5_weight_tid, %c2 : index + %q5_meta0 = index.assume %q5_meta0_0 [range(%q5_meta0_0, 0, 254)] : index + %q5_meta1_0 = index.add %q5_meta0, %c1 : index + %q5_meta1 = index.assume %q5_meta1_0 [range(%q5_meta1_0, 1, 255)] : index + view.store %q5_d0, %wsl_view[%q5_meta0] : f32, view<[%weight_metadata]xf32> + view.store %q5_d1, %wsl_view[%q5_meta1] : f32, view<[%weight_metadata]xf32> + view.store %q5_correction0, %wcl_view[%q5_meta0] : f32, view<[%weight_metadata]xf32> + view.store %q5_correction1, %wcl_view[%q5_meta1] : f32, view<[%weight_metadata]xf32> + + %q5_pair_word_base = index.mul %q5_pair, %c8 : index + %q5_group0_i32 = index.cast %q5_group0 : index to i32 + %q5_group1_i32 = index.cast %q5_group1 : index to i32 + %q5_group0_shift = vector.splat %q5_group0_i32 : vector<4xi32> + %q5_group1_shift = vector.splat %q5_group1_i32 : vector<4xi32> + %q5_nibble_mask = vector.constant 252645135 : vector<4xi32> + %q5_byte_ones = vector.constant 16843009 : vector<4xi32> + %q5_sign_extend = vector.constant 240 : vector<4xi32> + %q5_shift4 = vector.constant 4 : vector<4xi32> + %q5_weight_row_base = index.mul %q5_weight_tid, %k_wstride : index + %q5_w0 = index.assume %q5_weight_row_base [range(%q5_weight_row_base, 0, 10160)] : index + %q5_payload16 = index.constant 16 : index + %q5_payload48 = index.constant 48 : index + %q5_w0_hi0 = index.add %q5_weight_row_base, %q5_payload16 : index + %q5_w0_hi = index.assume %q5_w0_hi0 [range(%q5_w0_hi0, 16, 10176)] : index + %q5_w1_0 = index.add %q5_weight_row_base, %c32 : index + %q5_w1 = index.assume %q5_w1_0 [range(%q5_w1_0, 32, 10192)] : index + %q5_w1_hi0 = index.add %q5_weight_row_base, %q5_payload48 : index + %q5_w1_hi = index.assume %q5_w1_hi0 [range(%q5_w1_hi0, 48, 10208)] : index + + // First 16-byte payload half. + %q5_code_words0 = scf.if %is_q4_packed -> (vector<4xi32>) { + %field_add = index.constant 1 : index + %field_bytes = index.constant 1024 : offset + %field_pair = index.mul %q5_pair, %c2 : index + %field = index.add %field_pair, %field_add : index + %field_offset = index.scale %field, %field_bytes : index, offset -> offset + %payload_base = index.add %q5_block_byte_base, %field_offset : offset + %words = func.call @ggml_q8_1_x4_paired_weight_words(%paired_weights, %load_up, %src0_na, %peer_weight, %payload_base) : (i1, i1, buffer, buffer, offset) -> (vector<4xi32>) + scf.yield %words : vector<4xi32> + } else { + %words = vector.load %q5_code_view[%q5_pair_word_base] : view<32xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %q5_high_words0 = scf.if %is_q4 -> (vector<4xi32>) { + %zero = vector.constant 0 : vector<4xi32> + scf.yield %zero : vector<4xi32> + } else { + %high = vector.load %q5_high_view[%c0] : view<8xi32> -> vector<4xi32> + scf.yield %high : vector<4xi32> + } + %q5_low_codes0 = vector.andi %q5_code_words0, %q5_nibble_mask : vector<4xi32> + %q5_fifth00_shifted = vector.shrui %q5_high_words0, %q5_group0_shift : vector<4xi32> + %q5_fifth00_bits = vector.andi %q5_fifth00_shifted, %q5_byte_ones : vector<4xi32> + %q5_fifth00_unsigned = vector.shli %q5_fifth00_bits, %q5_shift4 : vector<4xi32> + %q5_fifth00_signed = vector.muli %q5_fifth00_bits, %q5_sign_extend : vector<4xi32> + %q5_fifth00 = scf.select %is_symi5, %q5_fifth00_signed, %q5_fifth00_unsigned : vector<4xi32> + %q5_values00_words = vector.ori %q5_low_codes0, %q5_fifth00 : vector<4xi32> + %q5_values00 = vector.bitcast %q5_values00_words : vector<4xi32> to vector<16xi8> + vector.store %q5_values00, %wl_flat[%q5_w0] : vector<16xi8>, view<[%weight_elements]xi8> + %q5_high_codes00 = vector.shrui %q5_code_words0, %q5_shift4 : vector<4xi32> + %q5_high_codes0 = vector.andi %q5_high_codes00, %q5_nibble_mask : vector<4xi32> + %q5_fifth10_shifted = vector.shrui %q5_high_words0, %q5_group1_shift : vector<4xi32> + %q5_fifth10_bits = vector.andi %q5_fifth10_shifted, %q5_byte_ones : vector<4xi32> + %q5_fifth10_unsigned = vector.shli %q5_fifth10_bits, %q5_shift4 : vector<4xi32> + %q5_fifth10_signed = vector.muli %q5_fifth10_bits, %q5_sign_extend : vector<4xi32> + %q5_fifth10 = scf.select %is_symi5, %q5_fifth10_signed, %q5_fifth10_unsigned : vector<4xi32> + %q5_values10_words = vector.ori %q5_high_codes0, %q5_fifth10 : vector<4xi32> + %q5_values10 = vector.bitcast %q5_values10_words : vector<4xi32> to vector<16xi8> + vector.store %q5_values10, %wl_flat[%q5_w1] : vector<16xi8>, view<[%weight_elements]xi8> + + // Second 16-byte payload half. Recompute its addresses after the first + // stores so only one half's decode payload remains live at a time. + %q5_pair_word_hi0 = index.add %q5_pair_word_base, %c4 : index + %q5_pair_word_hi = index.assume %q5_pair_word_hi0 [range(%q5_pair_word_hi0, 4, 28)] : index + %q5_code_words1 = scf.if %is_q4_packed -> (vector<4xi32>) { + %field_add = index.constant 2 : index + %field_bytes = index.constant 1024 : offset + %field_pair = index.mul %q5_pair, %c2 : index + %field = index.add %field_pair, %field_add : index + %field_offset = index.scale %field, %field_bytes : index, offset -> offset + %payload_base = index.add %q5_block_byte_base, %field_offset : offset + %words = func.call @ggml_q8_1_x4_paired_weight_words(%paired_weights, %load_up, %src0_na, %peer_weight, %payload_base) : (i1, i1, buffer, buffer, offset) -> (vector<4xi32>) + scf.yield %words : vector<4xi32> + } else { + %words = vector.load %q5_code_view[%q5_pair_word_hi] : view<32xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %q5_high_words1 = scf.if %is_q4 -> (vector<4xi32>) { + %zero = vector.constant 0 : vector<4xi32> + scf.yield %zero : vector<4xi32> + } else { + %high = vector.load %q5_high_view[%c4] : view<8xi32> -> vector<4xi32> + scf.yield %high : vector<4xi32> + } + %q5_low_codes1 = vector.andi %q5_code_words1, %q5_nibble_mask : vector<4xi32> + %q5_fifth01_shifted = vector.shrui %q5_high_words1, %q5_group0_shift : vector<4xi32> + %q5_fifth01_bits = vector.andi %q5_fifth01_shifted, %q5_byte_ones : vector<4xi32> + %q5_fifth01_unsigned = vector.shli %q5_fifth01_bits, %q5_shift4 : vector<4xi32> + %q5_fifth01_signed = vector.muli %q5_fifth01_bits, %q5_sign_extend : vector<4xi32> + %q5_fifth01 = scf.select %is_symi5, %q5_fifth01_signed, %q5_fifth01_unsigned : vector<4xi32> + %q5_values01_words = vector.ori %q5_low_codes1, %q5_fifth01 : vector<4xi32> + %q5_values01 = vector.bitcast %q5_values01_words : vector<4xi32> to vector<16xi8> + vector.store %q5_values01, %wl_flat[%q5_w0_hi] : vector<16xi8>, view<[%weight_elements]xi8> + %q5_high_codes10 = vector.shrui %q5_code_words1, %q5_shift4 : vector<4xi32> + %q5_high_codes1 = vector.andi %q5_high_codes10, %q5_nibble_mask : vector<4xi32> + %q5_fifth11_shifted = vector.shrui %q5_high_words1, %q5_group1_shift : vector<4xi32> + %q5_fifth11_bits = vector.andi %q5_fifth11_shifted, %q5_byte_ones : vector<4xi32> + %q5_fifth11_unsigned = vector.shli %q5_fifth11_bits, %q5_shift4 : vector<4xi32> + %q5_fifth11_signed = vector.muli %q5_fifth11_bits, %q5_sign_extend : vector<4xi32> + %q5_fifth11 = scf.select %is_symi5, %q5_fifth11_signed, %q5_fifth11_unsigned : vector<4xi32> + %q5_values11_words = vector.ori %q5_high_codes1, %q5_fifth11 : vector<4xi32> + %q5_values11 = vector.bitcast %q5_values11_words : vector<4xi32> to vector<16xi8> + vector.store %q5_values11, %wl_flat[%q5_w1_hi] : vector<16xi8>, view<[%weight_elements]xi8> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %step0_0, %step0_1, %step0_2, %step0_3, %step1_0, %step1_1, %step1_2, %step1_3 = scf.if %paired_weights -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %bk0 = index.mul %c0, %c32 : index + %bkh0 = index.add %bk0, %c16 : index + %iz0 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %lhs0_0_raw = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %bk0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lhs0_0_words = vector.bitcast %lhs0_0_raw : vector<16xi8> to vector<4xi32> + %lhs0_0 = vector.fragment %lhs0_0_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %lhs0_1_raw = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %bkh0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lhs0_1_words = vector.bitcast %lhs0_1_raw : vector<16xi8> to vector<4xi32> + %lhs0_1 = vector.fragment %lhs0_1_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %asi0_0 = index.add %asb0, %c0 : index + %asi0_1 = index.mul %asi0_0, %c2 : index + %asi0_2 = index.add %asi0_1, %lane_hi : index + %asi0_3 = index.mul %asi0_2, %c8 : index + %asv0, %auv0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %asi0_3) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + %n0_0 = index.constant 0 : index + %b0n0 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_0, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc0_0) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_1 = index.constant 16 : index + %b0n1 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_1, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc0_1) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_2 = index.constant 32 : index + %b0n2 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_2, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc0_2) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_3 = index.constant 48 : index + %b0n3 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_3, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc0_3) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_4 = index.constant 64 : index + %b0n4 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_4, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc1_0) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_5 = index.constant 80 : index + %b0n5 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_5, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc1_1) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_6 = index.constant 96 : index + %b0n6 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_6, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc1_2) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n0_7 = index.constant 112 : index + %b0n7 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c0, %n0_7, %lane_lo, %lhs0_0, %lhs0_1, %iz0, %asv0, %auv0, %fc1_3) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %bk1 = index.mul %c1, %c32 : index + %bkh1 = index.add %bk1, %c16 : index + %iz1 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %lhs1_0_raw = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %bk1, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lhs1_0_words = vector.bitcast %lhs1_0_raw : vector<16xi8> to vector<4xi32> + %lhs1_0 = vector.fragment %lhs1_0_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %lhs1_1_raw = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %bkh1, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lhs1_1_words = vector.bitcast %lhs1_1_raw : vector<16xi8> to vector<4xi32> + %lhs1_1 = vector.fragment %lhs1_1_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %asi1_0 = index.add %asb0, %c1 : index + %asi1_1 = index.mul %asi1_0, %c2 : index + %asi1_2 = index.add %asi1_1, %lane_hi : index + %asi1_3 = index.mul %asi1_2, %c8 : index + %asv1, %auv1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %asi1_3) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + %n1_0 = index.constant 0 : index + %b1n0 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_0, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n0) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_1 = index.constant 16 : index + %b1n1 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_1, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n1) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_2 = index.constant 32 : index + %b1n2 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_2, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n2) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_3 = index.constant 48 : index + %b1n3 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_3, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n3) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_4 = index.constant 64 : index + %b1n4 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_4, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n4) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_5 = index.constant 80 : index + %b1n5 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_5, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n5) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_6 = index.constant 96 : index + %b1n6 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_6, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n6) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + %n1_7 = index.constant 112 : index + %b1n7 = func.call @ggml_q8_1_x4_affine_rhs_tile(%scratch, %w_off, %ws_off, %wc_off, %c1, %n1_7, %lane_lo, %lhs1_0, %lhs1_1, %iz1, %asv1, %auv1, %b0n7) : (buffer, offset, offset, offset, index, index, index, vector<4xi32>, vector<4xi32>, vector<8xi32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) -> (vector<8xf32>) + scf.schedule.fence + scf.yield %b1n0, %b1n1, %b1n2, %b1n3, %b1n4, %b1n5, %b1n6, %b1n7 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } else { + // Two K32 blocks explicitly unrolled. Arithmetic order and the + // enclosing K64 staging/barrier cadence are unchanged. + %blk_k_u0 = index.mul %c0, %c32 : index + %blk_k16_u0 = index.add %blk_k_u0, %c16 : index + %iz_u0 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %wsi0_2_u0 = index.add %wsb0, %c0 : index + %wsi0_u0 = index.assume %wsi0_2_u0 [range(%wsi0_2_u0, 0, 127)] : index + %wsc0_u0 = view.load %wsl_view[%wsi0_u0] : view<[%weight_metadata]xf32> -> f32 + %wsv0_u0 = vector.splat %wsc0_u0 : vector<8xf32> + %wc0_u0 = view.load %wcl_view[%wsi0_u0] : view<[%weight_metadata]xf32> -> f32 + %wcv0_u0 = vector.splat %wc0_u0 : vector<8xf32> + %wsi1_2_u0 = index.add %wsb1, %c0 : index + %wsi1_u0 = index.assume %wsi1_2_u0 [range(%wsi1_2_u0, 0, 127)] : index + %wsc1_u0 = view.load %wsl_view[%wsi1_u0] : view<[%weight_metadata]xf32> -> f32 + %wsv1_u0 = vector.splat %wsc1_u0 : vector<8xf32> + %wc1_u0 = view.load %wcl_view[%wsi1_u0] : view<[%weight_metadata]xf32> -> f32 + %wcv1_u0 = vector.splat %wc1_u0 : vector<8xf32> + %wsi2_2_u0 = index.add %wsb2, %c0 : index + %wsi2_u0 = index.assume %wsi2_2_u0 [range(%wsi2_2_u0, 0, 127)] : index + %wsc2_u0 = view.load %wsl_view[%wsi2_u0] : view<[%weight_metadata]xf32> -> f32 + %wsv2_u0 = vector.splat %wsc2_u0 : vector<8xf32> + %wc2_u0 = view.load %wcl_view[%wsi2_u0] : view<[%weight_metadata]xf32> -> f32 + %wcv2_u0 = vector.splat %wc2_u0 : vector<8xf32> + %wsi3_2_u0 = index.add %wsb3, %c0 : index + %wsi3_u0 = index.assume %wsi3_2_u0 [range(%wsi3_2_u0, 0, 127)] : index + %wsc3_u0 = view.load %wsl_view[%wsi3_u0] : view<[%weight_metadata]xf32> -> f32 + %wsv3_u0 = vector.splat %wsc3_u0 : vector<8xf32> + %wc3_u0 = view.load %wcl_view[%wsi3_u0] : view<[%weight_metadata]xf32> -> f32 + %wcv3_u0 = vector.splat %wc3_u0 : vector<8xf32> + %lf0_0_raw_u0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %blk_k_u0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf0_0_words_u0 = vector.bitcast %lf0_0_raw_u0 : vector<16xi8> to vector<4xi32> + %lf0_0_u0 = vector.fragment %lf0_0_words_u0 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %lf0_1_raw_u0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %blk_k16_u0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf0_1_words_u0 = vector.bitcast %lf0_1_raw_u0 : vector<16xi8> to vector<4xi32> + %lf0_1_u0 = vector.fragment %lf0_1_words_u0 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %asi0_0_u0 = index.add %asb0, %c0 : index + %asi0_1_u0 = index.mul %asi0_0_u0, %c2 : index + %asi0_2_u0 = index.add %asi0_1_u0, %lane_hi : index + %asi0_3_u0 = index.mul %asi0_2_u0, %c8 : index + %asv0_u0, %auv0_u0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %asi0_3_u0) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + // Issue independent RHS fragments before consuming the first pair. + %rfs0_0_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_0_0_words_u0 = vector.bitcast %rfs0_0_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_0_0_u0 = vector.fragment %rfs0_0_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_0_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_0_1_words_u0 = vector.bitcast %rfs0_0_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_0_1_u0 = vector.fragment %rfs0_0_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_1_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_1_0_words_u0 = vector.bitcast %rfs0_1_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_1_0_u0 = vector.fragment %rfs0_1_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_1_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_1_1_words_u0 = vector.bitcast %rfs0_1_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_1_1_u0 = vector.fragment %rfs0_1_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_2_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_2_0_words_u0 = vector.bitcast %rfs0_2_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_2_0_u0 = vector.fragment %rfs0_2_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_2_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_2_1_words_u0 = vector.bitcast %rfs0_2_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_2_1_u0 = vector.fragment %rfs0_2_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_3_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_3_0_words_u0 = vector.bitcast %rfs0_3_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_3_0_u0 = vector.fragment %rfs0_3_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_3_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_3_1_words_u0 = vector.bitcast %rfs0_3_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs0_3_1_u0 = vector.fragment %rfs0_3_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %i0_0_0_u0 = vector.mma %lf0_0_u0, %rfs0_0_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_0_u0 = vector.mma %lf0_1_u0, %rfs0_0_1_u0, %i0_0_0_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_0_u0 = vector.sitofp %i1_0_0_u0 : vector<8xi32> to vector<8xf32> + %sv0_0_u0 = vector.mulf %asv0_u0, %wsv0_u0 : vector<8xf32> + %base_fn0_0_u0 = vector.fmaf %auv0_u0, %wcv0_u0, %fc0_0 : vector<8xf32> + %selected_base_fn0_0_u0 = scf.select %is_signed_rhs, %fc0_0, %base_fn0_0_u0 : vector<8xf32> + + %fn0_0_u0 = vector.fmaf %ff0_0_u0, %sv0_0_u0, %selected_base_fn0_0_u0 : vector<8xf32> + scf.schedule.fence + %i0_0_1_u0 = vector.mma %lf0_0_u0, %rfs0_1_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_1_u0 = vector.mma %lf0_1_u0, %rfs0_1_1_u0, %i0_0_1_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_1_u0 = vector.sitofp %i1_0_1_u0 : vector<8xi32> to vector<8xf32> + %sv0_1_u0 = vector.mulf %asv0_u0, %wsv1_u0 : vector<8xf32> + %base_fn0_1_u0 = vector.fmaf %auv0_u0, %wcv1_u0, %fc0_1 : vector<8xf32> + %selected_base_fn0_1_u0 = scf.select %is_signed_rhs, %fc0_1, %base_fn0_1_u0 : vector<8xf32> + + %fn0_1_u0 = vector.fmaf %ff0_1_u0, %sv0_1_u0, %selected_base_fn0_1_u0 : vector<8xf32> + scf.schedule.fence + %i0_0_2_u0 = vector.mma %lf0_0_u0, %rfs0_2_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_2_u0 = vector.mma %lf0_1_u0, %rfs0_2_1_u0, %i0_0_2_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_2_u0 = vector.sitofp %i1_0_2_u0 : vector<8xi32> to vector<8xf32> + %sv0_2_u0 = vector.mulf %asv0_u0, %wsv2_u0 : vector<8xf32> + %base_fn0_2_u0 = vector.fmaf %auv0_u0, %wcv2_u0, %fc0_2 : vector<8xf32> + %selected_base_fn0_2_u0 = scf.select %is_signed_rhs, %fc0_2, %base_fn0_2_u0 : vector<8xf32> + + %fn0_2_u0 = vector.fmaf %ff0_2_u0, %sv0_2_u0, %selected_base_fn0_2_u0 : vector<8xf32> + scf.schedule.fence + %i0_0_3_u0 = vector.mma %lf0_0_u0, %rfs0_3_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_3_u0 = vector.mma %lf0_1_u0, %rfs0_3_1_u0, %i0_0_3_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_3_u0 = vector.sitofp %i1_0_3_u0 : vector<8xi32> to vector<8xf32> + %sv0_3_u0 = vector.mulf %asv0_u0, %wsv3_u0 : vector<8xf32> + %base_fn0_3_u0 = vector.fmaf %auv0_u0, %wcv3_u0, %fc0_3 : vector<8xf32> + %selected_base_fn0_3_u0 = scf.select %is_signed_rhs, %fc0_3, %base_fn0_3_u0 : vector<8xf32> + + %fn0_3_u0 = vector.fmaf %ff0_3_u0, %sv0_3_u0, %selected_base_fn0_3_u0 : vector<8xf32> + scf.schedule.fence + %lf1_0_raw_u0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %blk_k_u0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf1_0_words_u0 = vector.bitcast %lf1_0_raw_u0 : vector<16xi8> to vector<4xi32> + %lf1_0_u0 = vector.fragment %lf1_0_words_u0 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %lf1_1_raw_u0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %blk_k16_u0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf1_1_words_u0 = vector.bitcast %lf1_1_raw_u0 : vector<16xi8> to vector<4xi32> + %lf1_1_u0 = vector.fragment %lf1_1_words_u0 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %asi1_0_u0 = index.add %asb1, %c0 : index + %asi1_1_u0 = index.mul %asi1_0_u0, %c2 : index + %asi1_2_u0 = index.add %asi1_1_u0, %lane_hi : index + %asi1_3_u0 = index.mul %asi1_2_u0, %c8 : index + %asv1_u0, %auv1_u0 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %asi1_3_u0) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + // Issue independent RHS fragments before consuming the first pair. + %rfs1_0_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_0_0_words_u0 = vector.bitcast %rfs1_0_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_0_0_u0 = vector.fragment %rfs1_0_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_0_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_0_1_words_u0 = vector.bitcast %rfs1_0_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_0_1_u0 = vector.fragment %rfs1_0_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_1_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_1_0_words_u0 = vector.bitcast %rfs1_1_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_1_0_u0 = vector.fragment %rfs1_1_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_1_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_1_1_words_u0 = vector.bitcast %rfs1_1_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_1_1_u0 = vector.fragment %rfs1_1_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_2_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_2_0_words_u0 = vector.bitcast %rfs1_2_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_2_0_u0 = vector.fragment %rfs1_2_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_2_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_2_1_words_u0 = vector.bitcast %rfs1_2_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_2_1_u0 = vector.fragment %rfs1_2_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_3_0_raw_u0 = vector.fragment.load %wl_view[%blk_k_u0, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_3_0_words_u0 = vector.bitcast %rfs1_3_0_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_3_0_u0 = vector.fragment %rfs1_3_0_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_3_1_raw_u0 = vector.fragment.load %wl_view[%blk_k16_u0, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_3_1_words_u0 = vector.bitcast %rfs1_3_1_raw_u0 : vector<16xi8> to vector<4xi32> + %rfs1_3_1_u0 = vector.fragment %rfs1_3_1_words_u0 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %i0_1_0_u0 = vector.mma %lf1_0_u0, %rfs1_0_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_0_u0 = vector.mma %lf1_1_u0, %rfs1_0_1_u0, %i0_1_0_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_0_u0 = vector.sitofp %i1_1_0_u0 : vector<8xi32> to vector<8xf32> + %sv1_0_u0 = vector.mulf %asv1_u0, %wsv0_u0 : vector<8xf32> + %base_fn1_0_u0 = vector.fmaf %auv1_u0, %wcv0_u0, %fc1_0 : vector<8xf32> + %selected_base_fn1_0_u0 = scf.select %is_signed_rhs, %fc1_0, %base_fn1_0_u0 : vector<8xf32> + + %fn1_0_u0 = vector.fmaf %ff1_0_u0, %sv1_0_u0, %selected_base_fn1_0_u0 : vector<8xf32> + scf.schedule.fence + %i0_1_1_u0 = vector.mma %lf1_0_u0, %rfs1_1_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_1_u0 = vector.mma %lf1_1_u0, %rfs1_1_1_u0, %i0_1_1_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_1_u0 = vector.sitofp %i1_1_1_u0 : vector<8xi32> to vector<8xf32> + %sv1_1_u0 = vector.mulf %asv1_u0, %wsv1_u0 : vector<8xf32> + %base_fn1_1_u0 = vector.fmaf %auv1_u0, %wcv1_u0, %fc1_1 : vector<8xf32> + %selected_base_fn1_1_u0 = scf.select %is_signed_rhs, %fc1_1, %base_fn1_1_u0 : vector<8xf32> + + %fn1_1_u0 = vector.fmaf %ff1_1_u0, %sv1_1_u0, %selected_base_fn1_1_u0 : vector<8xf32> + scf.schedule.fence + %i0_1_2_u0 = vector.mma %lf1_0_u0, %rfs1_2_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_2_u0 = vector.mma %lf1_1_u0, %rfs1_2_1_u0, %i0_1_2_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_2_u0 = vector.sitofp %i1_1_2_u0 : vector<8xi32> to vector<8xf32> + %sv1_2_u0 = vector.mulf %asv1_u0, %wsv2_u0 : vector<8xf32> + %base_fn1_2_u0 = vector.fmaf %auv1_u0, %wcv2_u0, %fc1_2 : vector<8xf32> + %selected_base_fn1_2_u0 = scf.select %is_signed_rhs, %fc1_2, %base_fn1_2_u0 : vector<8xf32> + + %fn1_2_u0 = vector.fmaf %ff1_2_u0, %sv1_2_u0, %selected_base_fn1_2_u0 : vector<8xf32> + scf.schedule.fence + %i0_1_3_u0 = vector.mma %lf1_0_u0, %rfs1_3_0_u0, %iz_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_3_u0 = vector.mma %lf1_1_u0, %rfs1_3_1_u0, %i0_1_3_u0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_3_u0 = vector.sitofp %i1_1_3_u0 : vector<8xi32> to vector<8xf32> + %sv1_3_u0 = vector.mulf %asv1_u0, %wsv3_u0 : vector<8xf32> + %base_fn1_3_u0 = vector.fmaf %auv1_u0, %wcv3_u0, %fc1_3 : vector<8xf32> + %selected_base_fn1_3_u0 = scf.select %is_signed_rhs, %fc1_3, %base_fn1_3_u0 : vector<8xf32> + + %fn1_3_u0 = vector.fmaf %ff1_3_u0, %sv1_3_u0, %selected_base_fn1_3_u0 : vector<8xf32> + scf.schedule.fence + %blk_k_u1 = index.mul %c1, %c32 : index + %blk_k16_u1 = index.add %blk_k_u1, %c16 : index + %iz_u1 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %wsi0_2_u1 = index.add %wsb0, %c1 : index + %wsi0_u1 = index.assume %wsi0_2_u1 [range(%wsi0_2_u1, 0, 127)] : index + %wsc0_u1 = view.load %wsl_view[%wsi0_u1] : view<[%weight_metadata]xf32> -> f32 + %wsv0_u1 = vector.splat %wsc0_u1 : vector<8xf32> + %wc0_u1 = view.load %wcl_view[%wsi0_u1] : view<[%weight_metadata]xf32> -> f32 + %wcv0_u1 = vector.splat %wc0_u1 : vector<8xf32> + %wsi1_2_u1 = index.add %wsb1, %c1 : index + %wsi1_u1 = index.assume %wsi1_2_u1 [range(%wsi1_2_u1, 0, 127)] : index + %wsc1_u1 = view.load %wsl_view[%wsi1_u1] : view<[%weight_metadata]xf32> -> f32 + %wsv1_u1 = vector.splat %wsc1_u1 : vector<8xf32> + %wc1_u1 = view.load %wcl_view[%wsi1_u1] : view<[%weight_metadata]xf32> -> f32 + %wcv1_u1 = vector.splat %wc1_u1 : vector<8xf32> + %wsi2_2_u1 = index.add %wsb2, %c1 : index + %wsi2_u1 = index.assume %wsi2_2_u1 [range(%wsi2_2_u1, 0, 127)] : index + %wsc2_u1 = view.load %wsl_view[%wsi2_u1] : view<[%weight_metadata]xf32> -> f32 + %wsv2_u1 = vector.splat %wsc2_u1 : vector<8xf32> + %wc2_u1 = view.load %wcl_view[%wsi2_u1] : view<[%weight_metadata]xf32> -> f32 + %wcv2_u1 = vector.splat %wc2_u1 : vector<8xf32> + %wsi3_2_u1 = index.add %wsb3, %c1 : index + %wsi3_u1 = index.assume %wsi3_2_u1 [range(%wsi3_2_u1, 0, 127)] : index + %wsc3_u1 = view.load %wsl_view[%wsi3_u1] : view<[%weight_metadata]xf32> -> f32 + %wsv3_u1 = vector.splat %wsc3_u1 : vector<8xf32> + %wc3_u1 = view.load %wcl_view[%wsi3_u1] : view<[%weight_metadata]xf32> -> f32 + %wcv3_u1 = vector.splat %wc3_u1 : vector<8xf32> + %lf0_0_raw_u1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %blk_k_u1, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf0_0_words_u1 = vector.bitcast %lf0_0_raw_u1 : vector<16xi8> to vector<4xi32> + %lf0_0_u1 = vector.fragment %lf0_0_words_u1 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %lf0_1_raw_u1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %blk_k16_u1, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf0_1_words_u1 = vector.bitcast %lf0_1_raw_u1 : vector<16xi8> to vector<4xi32> + %lf0_1_u1 = vector.fragment %lf0_1_words_u1 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %asi0_0_u1 = index.add %asb0, %c1 : index + %asi0_1_u1 = index.mul %asi0_0_u1, %c2 : index + %asi0_2_u1 = index.add %asi0_1_u1, %lane_hi : index + %asi0_3_u1 = index.mul %asi0_2_u1, %c8 : index + %asv0_u1, %auv0_u1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %asi0_3_u1) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + // Issue independent RHS fragments before consuming the first pair. + %rfs0_0_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_0_0_words_u1 = vector.bitcast %rfs0_0_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_0_0_u1 = vector.fragment %rfs0_0_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_0_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_0_1_words_u1 = vector.bitcast %rfs0_0_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_0_1_u1 = vector.fragment %rfs0_0_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_1_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_1_0_words_u1 = vector.bitcast %rfs0_1_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_1_0_u1 = vector.fragment %rfs0_1_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_1_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_1_1_words_u1 = vector.bitcast %rfs0_1_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_1_1_u1 = vector.fragment %rfs0_1_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_2_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_2_0_words_u1 = vector.bitcast %rfs0_2_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_2_0_u1 = vector.fragment %rfs0_2_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_2_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_2_1_words_u1 = vector.bitcast %rfs0_2_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_2_1_u1 = vector.fragment %rfs0_2_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_3_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_3_0_words_u1 = vector.bitcast %rfs0_3_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_3_0_u1 = vector.fragment %rfs0_3_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs0_3_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs0_3_1_words_u1 = vector.bitcast %rfs0_3_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs0_3_1_u1 = vector.fragment %rfs0_3_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %i0_0_0_u1 = vector.mma %lf0_0_u1, %rfs0_0_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_0_u1 = vector.mma %lf0_1_u1, %rfs0_0_1_u1, %i0_0_0_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_0_u1 = vector.sitofp %i1_0_0_u1 : vector<8xi32> to vector<8xf32> + %sv0_0_u1 = vector.mulf %asv0_u1, %wsv0_u1 : vector<8xf32> + %base_fn0_0 = vector.fmaf %auv0_u1, %wcv0_u1, %fn0_0_u0 : vector<8xf32> + %selected_base_fn0_0 = scf.select %is_signed_rhs, %fn0_0_u0, %base_fn0_0 : vector<8xf32> + + %fn0_0 = vector.fmaf %ff0_0_u1, %sv0_0_u1, %selected_base_fn0_0 : vector<8xf32> + scf.schedule.fence + %i0_0_1_u1 = vector.mma %lf0_0_u1, %rfs0_1_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_1_u1 = vector.mma %lf0_1_u1, %rfs0_1_1_u1, %i0_0_1_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_1_u1 = vector.sitofp %i1_0_1_u1 : vector<8xi32> to vector<8xf32> + %sv0_1_u1 = vector.mulf %asv0_u1, %wsv1_u1 : vector<8xf32> + %base_fn0_1 = vector.fmaf %auv0_u1, %wcv1_u1, %fn0_1_u0 : vector<8xf32> + %selected_base_fn0_1 = scf.select %is_signed_rhs, %fn0_1_u0, %base_fn0_1 : vector<8xf32> + + %fn0_1 = vector.fmaf %ff0_1_u1, %sv0_1_u1, %selected_base_fn0_1 : vector<8xf32> + scf.schedule.fence + %i0_0_2_u1 = vector.mma %lf0_0_u1, %rfs0_2_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_2_u1 = vector.mma %lf0_1_u1, %rfs0_2_1_u1, %i0_0_2_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_2_u1 = vector.sitofp %i1_0_2_u1 : vector<8xi32> to vector<8xf32> + %sv0_2_u1 = vector.mulf %asv0_u1, %wsv2_u1 : vector<8xf32> + %base_fn0_2 = vector.fmaf %auv0_u1, %wcv2_u1, %fn0_2_u0 : vector<8xf32> + %selected_base_fn0_2 = scf.select %is_signed_rhs, %fn0_2_u0, %base_fn0_2 : vector<8xf32> + + %fn0_2 = vector.fmaf %ff0_2_u1, %sv0_2_u1, %selected_base_fn0_2 : vector<8xf32> + scf.schedule.fence + %i0_0_3_u1 = vector.mma %lf0_0_u1, %rfs0_3_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_0_3_u1 = vector.mma %lf0_1_u1, %rfs0_3_1_u1, %i0_0_3_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff0_3_u1 = vector.sitofp %i1_0_3_u1 : vector<8xi32> to vector<8xf32> + %sv0_3_u1 = vector.mulf %asv0_u1, %wsv3_u1 : vector<8xf32> + %base_fn0_3 = vector.fmaf %auv0_u1, %wcv3_u1, %fn0_3_u0 : vector<8xf32> + %selected_base_fn0_3 = scf.select %is_signed_rhs, %fn0_3_u0, %base_fn0_3 : vector<8xf32> + + %fn0_3 = vector.fmaf %ff0_3_u1, %sv0_3_u1, %selected_base_fn0_3 : vector<8xf32> + scf.schedule.fence + %lf1_0_raw_u1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %blk_k_u1, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf1_0_words_u1 = vector.bitcast %lf1_0_raw_u1 : vector<16xi8> to vector<4xi32> + %lf1_0_u1 = vector.fragment %lf1_0_words_u1 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %lf1_1_raw_u1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %blk_k16_u1, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %lf1_1_words_u1 = vector.bitcast %lf1_1_raw_u1 : vector<16xi8> to vector<4xi32> + %lf1_1_u1 = vector.fragment %lf1_1_words_u1 shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %asi1_0_u1 = index.add %asb1, %c1 : index + %asi1_1_u1 = index.mul %asi1_0_u1, %c2 : index + %asi1_2_u1 = index.add %asi1_1_u1, %lane_hi : index + %asi1_3_u1 = index.mul %asi1_2_u1, %c8 : index + %asv1_u1, %auv1_u1 = func.call @ggml_q5_plane_dense_q8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %asi1_3_u1) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + // Issue independent RHS fragments before consuming the first pair. + %rfs1_0_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_0_0_words_u1 = vector.bitcast %rfs1_0_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_0_0_u1 = vector.fragment %rfs1_0_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_0_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln0] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_0_1_words_u1 = vector.bitcast %rfs1_0_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_0_1_u1 = vector.fragment %rfs1_0_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_1_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_1_0_words_u1 = vector.bitcast %rfs1_1_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_1_0_u1 = vector.fragment %rfs1_1_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_1_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln1] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_1_1_words_u1 = vector.bitcast %rfs1_1_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_1_1_u1 = vector.fragment %rfs1_1_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_2_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_2_0_words_u1 = vector.bitcast %rfs1_2_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_2_0_u1 = vector.fragment %rfs1_2_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_2_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln2] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_2_1_words_u1 = vector.bitcast %rfs1_2_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_2_1_u1 = vector.fragment %rfs1_2_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_3_0_raw_u1 = vector.fragment.load %wl_view[%blk_k_u1, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_3_0_words_u1 = vector.bitcast %rfs1_3_0_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_3_0_u1 = vector.fragment %rfs1_3_0_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %rfs1_3_1_raw_u1 = vector.fragment.load %wl_view[%blk_k16_u1, %ln3] shape [%c16, %c16] : view<64x[%k_n]xi8, %rhs_layout> -> vector<16xi8> + %rfs1_3_1_words_u1 = vector.bitcast %rfs1_3_1_raw_u1 : vector<16xi8> to vector<4xi32> + %rfs1_3_1_u1 = vector.fragment %rfs1_3_1_words_u1 shape [%c16, %c16] using {schema = %u8_schema : encoding} : vector<4xi32> + %i0_1_0_u1 = vector.mma %lf1_0_u1, %rfs1_0_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_0_u1 = vector.mma %lf1_1_u1, %rfs1_0_1_u1, %i0_1_0_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_0_u1 = vector.sitofp %i1_1_0_u1 : vector<8xi32> to vector<8xf32> + %sv1_0_u1 = vector.mulf %asv1_u1, %wsv0_u1 : vector<8xf32> + %base_fn1_0 = vector.fmaf %auv1_u1, %wcv0_u1, %fn1_0_u0 : vector<8xf32> + %selected_base_fn1_0 = scf.select %is_signed_rhs, %fn1_0_u0, %base_fn1_0 : vector<8xf32> + + %fn1_0 = vector.fmaf %ff1_0_u1, %sv1_0_u1, %selected_base_fn1_0 : vector<8xf32> + scf.schedule.fence + %i0_1_1_u1 = vector.mma %lf1_0_u1, %rfs1_1_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_1_u1 = vector.mma %lf1_1_u1, %rfs1_1_1_u1, %i0_1_1_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_1_u1 = vector.sitofp %i1_1_1_u1 : vector<8xi32> to vector<8xf32> + %sv1_1_u1 = vector.mulf %asv1_u1, %wsv1_u1 : vector<8xf32> + %base_fn1_1 = vector.fmaf %auv1_u1, %wcv1_u1, %fn1_1_u0 : vector<8xf32> + %selected_base_fn1_1 = scf.select %is_signed_rhs, %fn1_1_u0, %base_fn1_1 : vector<8xf32> + + %fn1_1 = vector.fmaf %ff1_1_u1, %sv1_1_u1, %selected_base_fn1_1 : vector<8xf32> + scf.schedule.fence + %i0_1_2_u1 = vector.mma %lf1_0_u1, %rfs1_2_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_2_u1 = vector.mma %lf1_1_u1, %rfs1_2_1_u1, %i0_1_2_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_2_u1 = vector.sitofp %i1_1_2_u1 : vector<8xi32> to vector<8xf32> + %sv1_2_u1 = vector.mulf %asv1_u1, %wsv2_u1 : vector<8xf32> + %base_fn1_2 = vector.fmaf %auv1_u1, %wcv2_u1, %fn1_2_u0 : vector<8xf32> + %selected_base_fn1_2 = scf.select %is_signed_rhs, %fn1_2_u0, %base_fn1_2 : vector<8xf32> + + %fn1_2 = vector.fmaf %ff1_2_u1, %sv1_2_u1, %selected_base_fn1_2 : vector<8xf32> + scf.schedule.fence + %i0_1_3_u1 = vector.mma %lf1_0_u1, %rfs1_3_0_u1, %iz_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i1_1_3_u1 = vector.mma %lf1_1_u1, %rfs1_3_1_u1, %i0_1_3_u1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %ff1_3_u1 = vector.sitofp %i1_1_3_u1 : vector<8xi32> to vector<8xf32> + %sv1_3_u1 = vector.mulf %asv1_u1, %wsv3_u1 : vector<8xf32> + %base_fn1_3 = vector.fmaf %auv1_u1, %wcv3_u1, %fn1_3_u0 : vector<8xf32> + %selected_base_fn1_3 = scf.select %is_signed_rhs, %fn1_3_u0, %base_fn1_3 : vector<8xf32> + + %fn1_3 = vector.fmaf %ff1_3_u1, %sv1_3_u1, %selected_base_fn1_3 : vector<8xf32> + scf.yield %fn0_0, %fn0_1, %fn0_2, %fn0_3, %fn1_0, %fn1_1, %fn1_2, %fn1_3 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %step0_0, %step0_1, %step0_2, %step0_3, %step1_0, %step1_1, %step1_2, %step1_3 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + + scf.if %paired_weights { + %silu0 = vector.siluf %f0_0 : vector<8xf32> + %result0 = vector.mulf %silu0, %f1_0 : vector<8xf32> + vector.fragment.store %result0, %dst_view[%gm0, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + %silu1 = vector.siluf %f0_1 : vector<8xf32> + %result1 = vector.mulf %silu1, %f1_1 : vector<8xf32> + vector.fragment.store %result1, %dst_view[%gm0, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + %silu2 = vector.siluf %f0_2 : vector<8xf32> + %result2 = vector.mulf %silu2, %f1_2 : vector<8xf32> + vector.fragment.store %result2, %dst_view[%gm0, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + %silu3 = vector.siluf %f0_3 : vector<8xf32> + %result3 = vector.mulf %silu3, %f1_3 : vector<8xf32> + vector.fragment.store %result3, %dst_view[%gm0, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + } else { + vector.fragment.store %f0_0, %dst_view[%gm0, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f0_1, %dst_view[%gm0, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f0_2, %dst_view[%gm0, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f0_3, %dst_view[%gm0, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f1_0, %dst_view[%gm1, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f1_1, %dst_view[%gm1, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f1_2, %dst_view[%gm1, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %f1_3, %dst_view[%gm1, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + } + template.return +} + +kernel.def target(@ggml_q5_plane_dense_gfx11_wave32) @ggml_mul_mat_q5_k_q8_plane_wmmai8_token256(%token_count: index) { + %unit = index.constant 1 : index + %tile_rows = index.constant 64 : index + %tile_tokens = index.constant 256 : index + %workgroup_size = index.constant 256 : index + %output_size = config.get @ggml.mul_mat_q5_k_q8_plane.output_size : index + %token_capacity = config.get @ggml.mul_mat_q5_k_q8_plane.token_capacity : index + %row_tiles = index.div %output_size, %tile_rows : index + %token_tiles = index.div %token_capacity, %tile_tokens : index + kernel.launch.config workgroups(%token_tiles, %row_tiles, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %base = index.constant 0 : offset + %lds_bytes = index.constant 30720 : offset + %w_off = index.constant 20480 : offset + %as_off = index.constant 25600 : offset + %ws_off = index.constant 27648 : offset + %asum_off = index.constant 28160 : offset + %wc_off = index.constant 30208 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 80 : index + %k_wstride = index.constant 80 : index + %k_m = index.constant 256 : index + %k_n = index.constant 64 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 8 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %is_iq4xs = scalar.constant false : i1 + %is_symi5 = scalar.constant true : i1 + %is_q4 = scalar.constant false : i1 + %is_q4_packed = scalar.constant false : i1 + %token256 = scalar.constant true : i1 + %q8_plane = scalar.constant true : i1 + %paired_weights = scalar.constant false : i1 + + %input_size0 = config.get @ggml.mul_mat_q5_k_q8_plane.input_size : index + %output_size0 = config.get @ggml.mul_mat_q5_k_q8_plane.output_size : index + %token_capacity = config.get @ggml.mul_mat_q5_k_q8_plane.token_capacity : index + %input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %token_capacity [range(%token_count, 256, 2048), mul(%token_count, 256), eq(%token_count, %token_capacity)] : index, index + %nchunks = index.div %input_size, %k_kc : index + %kblocks = index.div %input_size, %c32 : index + + %col_tile0 = kernel.workgroup.id : index + %row_tile0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 7)] : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 4095)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %wave = index.div %tid, %k_wave : index + + %weight_g = buffer.assume.memory_space %weight : buffer + %output_g = buffer.assume.memory_space %output : buffer + %q8_g = buffer.assume.memory_space %q8_input : buffer + %weight_noalias, %output_noalias, %q8_noalias = buffer.assume.noalias %weight_g, %output_g, %q8_g : buffer, buffer, buffer + + template.apply<@ggml.mul_mat_q5_k_q8_plane.wmmai8.body>(%is_iq4xs, %token256, %q8_plane, %paired_weights, %weight_noalias, %weight_noalias, %output_noalias, %q8_noalias, %base, %lds_bytes, %w_off, %as_off, %ws_off, %asum_off, %wc_off, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %input_size, %output_size, %bounded_token_count, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave, %is_symi5, %is_q4, %is_q4_packed) : (i1, i1, i1, i1, buffer, buffer, buffer, buffer, offset, offset, offset, offset, offset, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index, i1, i1, i1) + kernel.return +} + +// The JIT format and paired-projection constants remove inactive paths. +template.def<@ggml.mul_mat_q8_1_x4.entry> device @ggml_mul_mat_q8_1_x4_entry(%token_count: index, %paired_weights: i1, %q8_input: buffer, %weight: buffer, %peer_weight: buffer, %output: buffer) { + %base = index.constant 0 : offset + %lds_bytes_ordinary = index.constant 30720 : offset + %lds_bytes_paired = index.constant 24576 : offset + %lds_bytes = scf.select %paired_weights, %lds_bytes_paired, %lds_bytes_ordinary : offset + %w_off_ordinary = index.constant 20480 : offset + %w_off_paired = index.constant 10240 : offset + %w_off = scf.select %paired_weights, %w_off_paired, %w_off_ordinary : offset + %as_off_ordinary = index.constant 25600 : offset + %as_off_paired = index.constant 20480 : offset + %as_off = scf.select %paired_weights, %as_off_paired, %as_off_ordinary : offset + %ws_off_ordinary = index.constant 27648 : offset + %ws_off_paired = index.constant 21504 : offset + %ws_off = scf.select %paired_weights, %ws_off_paired, %ws_off_ordinary : offset + %asum_off_ordinary = index.constant 28160 : offset + %asum_off_paired = index.constant 22528 : offset + %asum_off = scf.select %paired_weights, %asum_off_paired, %asum_off_ordinary : offset + %wc_off_ordinary = index.constant 30208 : offset + %wc_off_paired = index.constant 23552 : offset + %wc_off = scf.select %paired_weights, %wc_off_paired, %wc_off_ordinary : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c23 = index.constant 23 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 80 : index + %k_wstride = index.constant 80 : index + %k_m_ordinary = index.constant 256 : index + %k_m_paired = index.constant 128 : index + %k_m = scf.select %paired_weights, %k_m_paired, %k_m_ordinary : index + %k_n_ordinary = index.constant 64 : index + %k_n_paired = index.constant 128 : index + %k_n = scf.select %paired_weights, %k_n_paired, %k_n_ordinary : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 8 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %weight_format = config.get @ggml.mul_mat_q8_1_x4.weight_format : index + %is_iq4xs = index.cmp eq, %weight_format, %c23 : index + %true = scalar.constant true : i1 + %token256 = scalar.xori %paired_weights, %true : i1 + %q8_plane = scalar.constant false : i1 + %is_symi5 = scalar.constant false : i1 + %q4_native = index.cmp eq, %weight_format, %c4 : index + %q4_packed_format = index.constant 44 : index + %is_q4_packed = index.cmp eq, %weight_format, %q4_packed_format : index + %is_q4 = scalar.ori %q4_native, %is_q4_packed : i1 + + %input_size0 = config.get @ggml.mul_mat_q8_1_x4.input_size : index + %output_size0 = config.get @ggml.mul_mat_q8_1_x4.output_size : index + %token_capacity = config.get @ggml.mul_mat_q8_1_x4.token_capacity : index + %input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %token_capacity [range(%token_count, 256, 2048), mul(%token_count, 256), eq(%token_count, %token_capacity)] : index, index + %nchunks = index.div %input_size, %k_kc : index + %kblocks = index.div %input_size, %c32 : index + + %col_tile0 = kernel.workgroup.id : index + %row_tile0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = scf.if %paired_weights -> (index) { + %col = index.assume %col_tile0 [range(%col_tile0, 0, 15)] : index + scf.yield %col : index + } else { + %col = index.assume %col_tile0 [range(%col_tile0, 0, 7)] : index + scf.yield %col : index + } + %row_tile = scf.if %paired_weights -> (index) { + %row = index.assume %row_tile0 [range(%row_tile0, 0, 8191)] : index + scf.yield %row : index + } else { + %row = index.assume %row_tile0 [range(%row_tile0, 0, 4095)] : index + scf.yield %row : index + } + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %wave = index.div %tid, %k_wave : index + + %peer_global = buffer.assume.memory_space %peer_weight : buffer + %weight_g = buffer.assume.memory_space %weight : buffer + %output_g = buffer.assume.memory_space %output : buffer + %q8_g = buffer.assume.memory_space %q8_input : buffer + %weight_noalias, %output_noalias, %q8_noalias = buffer.assume.noalias %weight_g, %output_g, %q8_g : buffer, buffer, buffer + + template.apply<@ggml.mul_mat_q5_k_q8_plane.wmmai8.body>(%is_iq4xs, %token256, %q8_plane, %paired_weights, %peer_global, %weight_noalias, %output_noalias, %q8_noalias, %base, %lds_bytes, %w_off, %as_off, %ws_off, %asum_off, %wc_off, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %input_size, %output_size, %bounded_token_count, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave, %is_symi5, %is_q4, %is_q4_packed) : (i1, i1, i1, i1, buffer, buffer, buffer, buffer, offset, offset, offset, offset, offset, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index, i1, i1, i1) + template.return +} + +kernel.def target(@ggml_q5_plane_dense_gfx11_wave32) export("ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256") @ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256(%token_count: index) { + %unit = index.constant 1 : index + %tile_rows = index.constant 64 : index + %tile_tokens = index.constant 256 : index + %workgroup_size = index.constant 256 : index + %output_size = config.get @ggml.mul_mat_q8_1_x4.output_size : index + %token_capacity = config.get @ggml.mul_mat_q8_1_x4.token_capacity : index + %row_tiles = index.div %output_size, %tile_rows : index + %token_tiles = index.div %token_capacity, %tile_tokens : index + kernel.launch.config workgroups(%token_tiles, %row_tiles, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %paired_weights = scalar.constant false : i1 + template.apply<@ggml.mul_mat_q8_1_x4.entry>(%token_count, %paired_weights, %q8_input, %weight, %weight, %output) : (index, i1, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_q5_plane_dense_gfx11_wave32) export("ggml_mul_mat_q4_k_q8_1_x4_swiglu_wmma_token256") @ggml_mul_mat_q4_k_q8_1_x4_swiglu_wmma_token256(%token_count: index) { + %unit = index.constant 1 : index + %tile_rows = index.constant 64 : index + %tile_tokens = index.constant 128 : index + %workgroup_size = index.constant 256 : index + %output_size = config.get @ggml.mul_mat_q8_1_x4.output_size : index + %token_capacity = config.get @ggml.mul_mat_q8_1_x4.token_capacity : index + %row_tiles = index.div %output_size, %tile_rows : index + %token_tiles = index.div %token_capacity, %tile_tokens : index + kernel.launch.config workgroups(%token_tiles, %row_tiles, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer) { + %paired_weights = scalar.constant true : i1 + template.apply<@ggml.mul_mat_q8_1_x4.entry>(%token_count, %paired_weights, %q8_input, %weight, %up_weight, %output) : (index, i1, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom new file mode 100644 index 000000000000..52341d38eeb9 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom @@ -0,0 +1,1580 @@ +func.decl @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer) + +config.decl @ggml.mul_mat_q6_k_packed.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_q6_k_packed.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.mul_mat_q6_k_packed.output_accumulation : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.mul_mat_q6_k_packed.weight_offset : %value: index where [range(%value, 0, 2147483647), mul(%value, 256)] + +config.decl @ggml.mul_mat_q6_k_packed.token_capacity : %value: index where [range(%value, 128, 2048), mul(%value, 128)] + +config.decl @ggml.mul_mat_q6_k_shortlist.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +config.decl @ggml.mul_mat_q6_k_shortlist.output_size : %value: index where [range(%value, 64, 262144), mul(%value, 64)] + +config.decl @ggml.mul_mat_q6_k_shortlist.selected_group_count : %value: index where [range(%value, 1, 1024)] + +template.decl @ggml.mul_mat_q6_k_packed.publish_f16_wave32(%result: vector<16xf16>, %token_offset: index, %token_tile_base: index, %bounded_token_count: index, %bounded_output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) + +template.decl @ggml.mul_mat_q6_k_packed.publish_token1_overwrite_f32_wave32(%result: vector<8xf32>, %output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %output: buffer) + +template.decl @ggml.mul_mat_q6_k_packed.token1.body(%token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_q6_k_i8_prepacked.body(%token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %input: buffer, %weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_q6_k_packed.selected_refine.body(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %candidate_count0: index, %input: buffer, %weight: buffer, %candidates: buffer, %output: buffer) + +template.decl @ggml.mul_mat_q6_k_packed.prefill_wave32.body(%input_is_f16: i1, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %output: buffer) + +amdgpu.target @ggml_mul_mat_q6_k_prefill_gfx11_wave32 {subgroup_size = 32} + +func.decl @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) + +template.def<@ggml.mul_mat_q6_k_packed.publish_f16_wave32> device @ggml_mul_mat_q6_k_packed_publish_f16_wave32(%result: vector<16xf16>, %token_offset: index, %token_tile_base: index, %bounded_token_count: index, %bounded_output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) { + %lane = kernel.subgroup.lane.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c0_offset = index.constant 0 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %c0_f32x8 = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %publish_token0 = index.div %lane, %c2 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c2 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index + %publish_channel_add = index.mul %publish_packet, %c8 : index + %token_base = index.add %token_tile_base, %token_offset : index + %token = index.add %token_base, %publish_token : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel = index.add %subgroup_channel_base, %publish_channel_add : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %writes = scalar.andi %valid_token, %valid_channel : i1 + %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index + vector.fragment.store %result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes { + %bounded_token, %output_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %bounded_output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_token_count]x[%bounded_output_size]xf32> + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> + %wide = vector.extf %values : vector<8xf16> to vector<8xf32> + %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> + %published = scf.if %accumulates_output -> (vector<8xf32>) { + %residual = vector.load.mask %bounded_output_view[%bounded_token, %channel], %mask, %c0_f32x8 : view<[%output_token_count]x[%bounded_output_size]xf32>, vector<8xi1>, vector<8xf32> + %sum = vector.addf %residual, %wide : vector<8xf32> + scf.yield %sum : vector<8xf32> + } else { + scf.yield %wide : vector<8xf32> + } + vector.store.mask %published, %bounded_output_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%output_token_count]x[%bounded_output_size]xf32>, vector<8xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + template.return +} + +template.def<@ggml.mul_mat_q6_k_packed.token1.body> device @ggml_mul_mat_q6_k_packed_token1_body(%token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %scale_row_layout: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c13440_bytes = index.constant 13440 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %empty_words = vector.constant 0 : vector<4xi32> + %zero_scale_bytes = vector.constant 0 : vector<8xi8> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high0_mask = vector.constant 50529027 : vector<4xi32> + %high1_mask = vector.constant 202116108 : vector<4xi32> + %high2_mask = vector.constant 808464432 : vector<4xi32> + %high3_mask = vector.constant -1061109568 : vector<4xi32> + %shift2 = vector.constant 2 : vector<4xi32> + %shift4 = vector.constant 4 : vector<4xi32> + %negative32_f32 = vector.constant -32.0 : vector<16xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x16xf16, %activation_fragment_layout> + + %channel_tile_base = index.mul %channel_tile, %c64 : index + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %channel = index.add %channel_tile_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %use_scale_row_layout = index.cmp eq, %scale_row_layout, %c1 : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<8xf32>) -> (vector<8xf32>) { + %row_group_block0 = index.mul %channel_tile, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %tile_byte_base = index.scale %row_group_block, %c13440_bytes : index, offset -> offset + %scale_byte_base = index.add %tile_byte_base, %c128_bytes : offset + %raw_byte_base = index.add %tile_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%tile_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %scale_row_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_noalias[%raw_byte_base] : buffer -> view<2x3x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %load_scale_row = scalar.andi %valid_channel, %use_scale_row_layout : i1 + %scale_bytes = scf.if %load_scale_row -> (vector<8xi8>) { + %loaded_scale_bytes = vector.load %scale_row_view[%weight_local_row, %weight_half, %c0] : view<64x2x8xi8> -> vector<8xi8> + scf.yield %loaded_scale_bytes : vector<8xi8> + } else { + scf.yield %zero_scale_bytes : vector<8xi8> + } + %block_result, %final_ql0_words, %final_ql1_words, %final_qh_words = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<8xf32>, %cached_ql0_words = %empty_words : vector<4xi32>, %cached_ql1_words = %empty_words : vector<4xi32>, %cached_qh_words = %empty_words : vector<4xi32>) -> (vector<8xf32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %quant_group_next = index.add %quant_group, %c1 : index + %half128 = index.div %quant_group, %c4 : index + %group_in_half = index.rem %quant_group, %c4 : index + %is_low_pair = index.cmp eq, %group_in_half, %c0 : index + %is_first_group = index.cmp eq, %quant_group, %c0 : index + %prefetch_second_half = index.cmp eq, %quant_group, %c2 : index + %weight_values0, %weight_values1, %next_ql0_words, %next_ql1_words, %next_qh_words = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_words, %ql1_words, %qh_words = scf.if %is_first_group -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_bytes = vector.load %raw_view[%half128, %c0, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %ql1_bytes = vector.load %raw_view[%half128, %c1, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %qh_bytes = vector.load %raw_view[%half128, %c2, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %loaded_ql0_words = vector.bitcast %ql0_bytes : vector<16xi8> to vector<4xi32> + %loaded_ql1_words = vector.bitcast %ql1_bytes : vector<16xi8> to vector<4xi32> + %loaded_qh_words = vector.bitcast %qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %loaded_ql0_words, %loaded_ql1_words, %loaded_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %codes0_words, %codes1_words = scf.if %is_low_pair -> (vector<4xi32>, vector<4xi32>) { + %ql0_low = vector.andi %ql0_words, %nibble_mask : vector<4xi32> + %ql1_low = vector.andi %ql1_words, %nibble_mask : vector<4xi32> + %qh0_bits0 = vector.andi %qh_words, %high0_mask : vector<4xi32> + %qh1_bits0 = vector.andi %qh_words, %high1_mask : vector<4xi32> + %qh0_bits = vector.shli %qh0_bits0, %shift4 : vector<4xi32> + %qh1_bits = vector.shli %qh1_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_low, %qh0_bits : vector<4xi32> + %codes1 = vector.ori %ql1_low, %qh1_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } else { + %ql0_high0 = vector.shrui %ql0_words, %shift4 : vector<4xi32> + %ql1_high0 = vector.shrui %ql1_words, %shift4 : vector<4xi32> + %ql0_high = vector.andi %ql0_high0, %nibble_mask : vector<4xi32> + %ql1_high = vector.andi %ql1_high0, %nibble_mask : vector<4xi32> + %qh2_bits = vector.andi %qh_words, %high2_mask : vector<4xi32> + %qh3_bits0 = vector.andi %qh_words, %high3_mask : vector<4xi32> + %qh3_bits = vector.shrui %qh3_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_high, %qh2_bits : vector<4xi32> + %codes1 = vector.ori %ql1_high, %qh3_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } + %codes0_u8 = vector.bitcast %codes0_words : vector<4xi32> to vector<16xi8> + %codes1_u8 = vector.bitcast %codes1_words : vector<4xi32> to vector<16xi8> + %codes0_unsigned_f32 = vector.uitofp %codes0_u8 : vector<16xi8> to vector<16xf32> + %codes1_unsigned_f32 = vector.uitofp %codes1_u8 : vector<16xi8> to vector<16xf32> + %codes0_f32 = vector.addf %codes0_unsigned_f32, %negative32_f32 : vector<16xf32> + %codes1_f32 = vector.addf %codes1_unsigned_f32, %negative32_f32 : vector<16xf32> + %scale0_i8, %scale1_i8 = scf.if %use_scale_row_layout -> (i8, i8) { + %row_scale0_i8 = vector.extract %scale_bytes[%quant_group] : vector<8xi8> -> i8 + %row_scale1_i8 = vector.extract %scale_bytes[%quant_group_next] : vector<8xi8> -> i8 + scf.yield %row_scale0_i8, %row_scale1_i8 : i8, i8 + } else { + %group_scale0_i8 = view.load %scale_view[%quant_group, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %group_scale1_i8 = view.load %scale_view[%quant_group_next, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + scf.yield %group_scale0_i8, %group_scale1_i8 : i8, i8 + } + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words = scf.if %prefetch_second_half -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %next_ql0_bytes = vector.load %raw_view[%c1, %c0, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql1_bytes = vector.load %raw_view[%c1, %c1, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_qh_bytes = vector.load %raw_view[%c1, %c2, %weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql0 = vector.bitcast %next_ql0_bytes : vector<16xi8> to vector<4xi32> + %next_ql1 = vector.bitcast %next_ql1_bytes : vector<16xi8> to vector<4xi32> + %next_qh = vector.bitcast %next_qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %next_ql0, %next_ql1, %next_qh : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %ql0_words, %ql1_words, %qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %values0, %values1, %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1, %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<8xf32>) -> (vector<8xf32>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs, %rhs, %half_acc : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.yield %half_next : vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next, %next_ql0_words, %next_ql1_words, %next_qh_words : vector<8xf32>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %block_result : vector<8xf32> + } + + %overwrites_output = index.cmp eq, %output_accumulation, %c0 : index + %has_one_token = index.cmp eq, %bounded_token_count, %c1 : index + %use_token1_overwrite = scalar.andi %overwrites_output, %has_one_token : i1 + scf.if %use_token1_overwrite { + template.apply<@ggml.mul_mat_q6_k_packed.publish_token1_overwrite_f32_wave32>(%result, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias) : (vector<8xf32>, index, index, index, buffer) + } else { + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + } + template.return +} + +template.def<@ggml.mul_mat_q6_k_packed.publish_token1_overwrite_f32_wave32> device @ggml_mul_mat_q6_k_packed_publish_token1_overwrite_f32_wave32(%result: vector<8xf32>, %output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %output: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %channel = index.add %channel_tile_base, %subgroup_channel_add : index + %lane = kernel.subgroup.lane.id : index + %lane_token = index.rem %lane, %c16 : index + %lane_channel = index.div %lane, %c16 : index + %partial_channel_base = index.add %channel, %lane_channel : index + %stores_token0 = index.cmp eq, %lane_token, %c0 : index + scf.if %stores_token0 { + %dst = buffer.view %output[%base] : buffer -> view<1x[%output_size]xf32> + scf.for %element = [%c0 to %c8 step %c1] unroll { + %channel_step = index.mul %element, %c2 : index + %partial_channel = index.add %partial_channel_base, %channel_step : index + %channel_valid = index.cmp ult, %partial_channel, %output_size : index + scf.if %channel_valid { + %bounded_channel, %output_channels = index.assume %partial_channel, %output_size [lt(%partial_channel, %output_size)] : index, index + %value = vector.extract %result[%element] : vector<8xf32> -> f32 + view.store %value, %dst[%c0, %bounded_channel] : f32, view<1x[%output_size]xf32> + } + } + } + template.return +} + +kernel.def export("ggml_mul_mat_q6_k_packed_token1_f16_wmma") @ggml_mul_mat_q6_k_packed_token1_f16_wmma(%token_count: index) { + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_q6_k_packed.input_size : index + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %scale_row_layout = index.constant 1 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %channel_tile = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_q6_k_packed.token1.body>(%bounded_token_count, %input_size, %output_size, %channel_tile, %scale_row_layout, %input, %weight, %output) : (index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +template.def<@ggml.mul_mat_q6_k_i8_prepacked.body> device @ggml_mul_mat_q6_k_i8_prepacked_body(%token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c17536_bytes = index.constant 17536 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x16xf16, %activation_fragment_layout> + + %channel_tile_base = index.mul %channel_tile, %c64 : index + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %channel = index.add %channel_tile_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<8xf32>) -> (vector<8xf32>) { + %row_group_block0 = index.mul %channel_tile, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %tile_byte_base = index.scale %row_group_block, %c17536_bytes : index, offset -> offset + %scale_byte_base = index.add %tile_byte_base, %c128_bytes : offset + %code_byte_base = index.add %tile_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%tile_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %code_view = buffer.view %weight_noalias[%code_byte_base] : buffer -> view<8x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %c0_f32 = scalar.constant 0.0 : f32 + scf.yield %c0_f32 : f32 + } + %block_result = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<8xf32>) -> (vector<8xf32>) { + %quant_group_next = index.add %quant_group, %c1 : index + %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { + %codes0 = vector.load %code_view[%quant_group, %weight_local_row, %weight_k] : view<8x64x32xi8> -> vector<16xi8> + %codes1 = vector.load %code_view[%quant_group_next, %weight_local_row, %weight_k] : view<8x64x32xi8> -> vector<16xi8> + %codes0_f32 = vector.sitofp %codes0 : vector<16xi8> to vector<16xf32> + %codes1_f32 = vector.sitofp %codes1 : vector<16xi8> to vector<16xf32> + %scale0_i8 = view.load %scale_view[%quant_group, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %scale1_i8 = view.load %scale_view[%quant_group_next, %weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + scf.yield %values0, %values1 : vector<16xf16>, vector<16xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1 : vector<16xf16>, vector<16xf16> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<8xf32>) -> (vector<8xf32>) unroll { + %lhs_fragment = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs_fragment = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs_fragment, %rhs_fragment, %half_acc : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.yield %half_next : vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next : vector<8xf32> + } + scf.yield %block_result : vector<8xf32> + } + + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.return +} + +kernel.def export("ggml_mul_mat_q6_k_i8_prepacked_f16_wmma") @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma(%token_count: index) { + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %c1 = index.constant 1 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_q6_k_packed.input_size : index + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %channel_tile = kernel.workgroup.id : index + template.apply<@ggml.mul_mat_q6_k_i8_prepacked.body>(%bounded_token_count, %input_size, %output_size, %channel_tile, %input, %weight, %output) : (index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +func.def inline @ggml_expand_symmetric_i2_words(%q2_packed: vector<2xi32>) -> (vector<4xi32>) { + %q2_mask16 = vector.constant 65535 : vector<2xi32> + %q2_mask8 = vector.constant 16711935 : vector<2xi32> + %q2_mask4 = vector.constant 252645135 : vector<2xi32> + %q2_mask2 = vector.constant 858993459 : vector<2xi32> + %q2_sign = vector.constant 572662306 : vector<2xi32> + %q2_shift16 = vector.constant 16 : vector<2xi32> + %q2_shift8 = vector.constant 8 : vector<2xi32> + %q2_shift4 = vector.constant 4 : vector<2xi32> + %q2_shift2 = vector.constant 2 : vector<2xi32> + %q2_shift1 = vector.constant 1 : vector<2xi32> + %q2_lo0 = vector.andi %q2_packed, %q2_mask16 : vector<2xi32> + %q2_hi0 = vector.shrui %q2_packed, %q2_shift16 : vector<2xi32> + %q2_lo1s = vector.shli %q2_lo0, %q2_shift8 : vector<2xi32> + %q2_lo1o = vector.ori %q2_lo0, %q2_lo1s : vector<2xi32> + %q2_lo1 = vector.andi %q2_lo1o, %q2_mask8 : vector<2xi32> + %q2_lo2s = vector.shli %q2_lo1, %q2_shift4 : vector<2xi32> + %q2_lo2o = vector.ori %q2_lo1, %q2_lo2s : vector<2xi32> + %q2_lo2 = vector.andi %q2_lo2o, %q2_mask4 : vector<2xi32> + %q2_lo3s = vector.shli %q2_lo2, %q2_shift2 : vector<2xi32> + %q2_lo3o = vector.ori %q2_lo2, %q2_lo3s : vector<2xi32> + %q2_lo3 = vector.andi %q2_lo3o, %q2_mask2 : vector<2xi32> + %q2_losign = vector.andi %q2_lo3, %q2_sign : vector<2xi32> + %q2_losign2 = vector.shli %q2_losign, %q2_shift1 : vector<2xi32> + %q2_losign3 = vector.shli %q2_losign, %q2_shift2 : vector<2xi32> + %q2_lo4 = vector.ori %q2_lo3, %q2_losign2 : vector<2xi32> + %q2_lo = vector.ori %q2_lo4, %q2_losign3 : vector<2xi32> + %q2_hi1s = vector.shli %q2_hi0, %q2_shift8 : vector<2xi32> + %q2_hi1o = vector.ori %q2_hi0, %q2_hi1s : vector<2xi32> + %q2_hi1 = vector.andi %q2_hi1o, %q2_mask8 : vector<2xi32> + %q2_hi2s = vector.shli %q2_hi1, %q2_shift4 : vector<2xi32> + %q2_hi2o = vector.ori %q2_hi1, %q2_hi2s : vector<2xi32> + %q2_hi2 = vector.andi %q2_hi2o, %q2_mask4 : vector<2xi32> + %q2_hi3s = vector.shli %q2_hi2, %q2_shift2 : vector<2xi32> + %q2_hi3o = vector.ori %q2_hi2, %q2_hi3s : vector<2xi32> + %q2_hi3 = vector.andi %q2_hi3o, %q2_mask2 : vector<2xi32> + %q2_hisign = vector.andi %q2_hi3, %q2_sign : vector<2xi32> + %q2_hisign2 = vector.shli %q2_hisign, %q2_shift1 : vector<2xi32> + %q2_hisign3 = vector.shli %q2_hisign, %q2_shift2 : vector<2xi32> + %q2_hi4 = vector.ori %q2_hi3, %q2_hisign2 : vector<2xi32> + %q2_hi = vector.ori %q2_hi4, %q2_hisign3 : vector<2xi32> + %weight_word0 = vector.extract %q2_lo[0] : vector<2xi32> -> i32 + %weight_word1 = vector.extract %q2_hi[0] : vector<2xi32> -> i32 + %weight_word2 = vector.extract %q2_lo[1] : vector<2xi32> -> i32 + %weight_word3 = vector.extract %q2_hi[1] : vector<2xi32> -> i32 + %weight_words = vector.from_elements %weight_word0, %weight_word1, %weight_word2, %weight_word3 : vector<4xi32> + func.return %weight_words : vector<4xi32> +} + +func.def inline @ggml_symmetric_i2_dot_preloaded(%acc: f32, %weight_words: vector<4xi32>, %activation_words: vector<4xi32>, %combined_scale: f32) -> (f32) { + %zero_i32 = scalar.constant 0 : i32 + %zero_i32x4 = vector.constant 0 : vector<4xi32> + %dot_parts = vector.dot8i4 %weight_words, %activation_words, %zero_i32x4 : vector<4xi32> + %dot_i32 = vector.reduce %dot_parts, %zero_i32 : vector<4xi32>, i32 + %dot = scalar.sitofp %dot_i32 : i32 to f32 + %next = scalar.fmaf %dot, %combined_scale, %acc : f32 + func.return %next : f32 +} + +kernel.def export("ggml_select_symmetric_i4_k32_groups") @ggml_select_symmetric_i4_k32_groups() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%input: buffer, %selected_groups: buffer) { + %input_size0 = config.get @ggml.mul_mat_q6_k_shortlist.input_size : index + %input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %selected_count = config.get @ggml.mul_mat_q6_k_shortlist.selected_group_count : index + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c32 = index.constant 32 : index + %base = index.constant 0 : offset + %zero_f32 = scalar.constant 0.0 : f32 + %negative_large = scalar.constant -3.4028234663852886e+38 : f32 + %zero_i32 = scalar.constant 0 : i32 + %one_i32 = scalar.constant 1 : i32 + %largest_i32 = scalar.constant 2147483647 : i32 + %group_count0 = index.div %input_size, %c32 : index + %group_count = index.assume %group_count0 [range(%group_count0, 8, 256)] : index + %valid = index.cmp ult, %tid, %group_count : index + %input_global = buffer.assume.memory_space %input : buffer + %groups_global = buffer.assume.memory_space %selected_groups : buffer + %input_noalias, %groups_noalias = buffer.assume.noalias %input_global, %groups_global : buffer, buffer + %input_view = buffer.view %input_noalias[%base] : buffer -> view<[%input_size]xf32> + %groups_view = buffer.view %groups_noalias[%base] : buffer -> view<[%selected_count]xi32> + %score = scf.if %valid -> (f32) { + %group_base = index.mul %tid, %c32 : index + %sum = scf.for %offset = [%c0 to %c32 step %c4](%iter = %zero_f32 : f32) -> (f32) unroll { + %element = index.add %group_base, %offset : index + %values = vector.load %input_view[%element] : view<[%input_size]xf32> -> vector<4xf32> + %absolute = vector.absf %values : vector<4xf32> + %partial = vector.reduce %absolute, %zero_f32 : vector<4xf32>, f32 + %next = scalar.addf %iter, %partial : f32 + scf.yield %next : f32 + } + scf.yield %sum : f32 + } else { + scf.yield %negative_large : f32 + } + %tid_i32 = index.cast %tid : index to i32 + %final = scf.for %rank = [%c0 to %selected_count step %c1](%remaining = %score : f32) -> (f32) { + %winner_score = kernel.workgroup.reduce %remaining : f32 + %matches = scalar.cmpf oeq, %remaining, %winner_score : f32 + %winner_candidate = scf.select %matches, %tid_i32, %largest_i32 : i32 + %winner_id = kernel.workgroup.reduce %winner_candidate : i32 + %is_winner = scalar.cmpi eq, %tid_i32, %winner_id : i32 + %next = scf.select %is_winner, %negative_large, %remaining : f32 + scf.yield %next : f32 + } + %matches_selected_sentinel = scalar.cmpf oeq, %final, %negative_large : f32 + %is_selected = scalar.andi %valid, %matches_selected_sentinel : i1 + %selected_i32 = scf.select %is_selected, %one_i32, %zero_i32 : i32 + %rank_i32 = kernel.workgroup.scan %selected_i32 {direction = forward, mode = exclusive} : i32 + scf.if %is_selected { + %rank0 = index.cast %rank_i32 : i32 to index + %bounded_rank = index.assume %rank0 [range(%rank0, 0, 255)] : index + %rank, %bounded_selected_count = index.assume %bounded_rank, %selected_count [lt(%bounded_rank, %selected_count)] : index, index + view.store %tid_i32, %groups_view[%rank] : i32, view<[%selected_count]xi32> + } + kernel.return +} + +kernel.def export("ggml_mul_mat_q6_k_symmetric_i2_scan_token1") @ggml_mul_mat_q6_k_symmetric_i2_scan_token1() { + %rows = config.get @ggml.mul_mat_q6_k_shortlist.output_size : index + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + %c511 = index.constant 511 : index + %c512 = index.constant 512 : index + %rows_up = index.add %rows, %c511 : index + %workgroups = index.div %rows_up, %c512 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%weight: buffer, %output: buffer, %qact: buffer, %scales: buffer, %selected_groups: buffer) { + %input_size0 = config.get @ggml.mul_mat_q6_k_shortlist.input_size : index + %output_size0 = config.get @ggml.mul_mat_q6_k_shortlist.output_size : index + %selected_count0 = config.get @ggml.mul_mat_q6_k_shortlist.selected_group_count : index + %input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %selected_count = index.assume %selected_count0 [range(%selected_count0, 1, 1024)] : index + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + %workgroup = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %cohort_base = index.mul %workgroup, %c128 : index + %cohort = index.add %cohort_base, %workitem : index + %row = index.mul %cohort, %c4 : index + %valid_cohort = index.cmp ult, %row, %output_size : index + %block_count = index.div %input_size, %c256 : index + %group_count = index.div %input_size, %c32 : index + %rows_plus = index.add %output_size, %c255 : index + %row_group_units = index.div %rows_plus, %c256 : index + %row_group = index.mul %row_group_units, %c32 : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c8 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + %weight_global = buffer.assume.memory_space %weight : buffer + %output_global = buffer.assume.memory_space %output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %groups_global = buffer.assume.memory_space %selected_groups : buffer + %weight_na, %output_na, %qact_na, %scales_na, %groups_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global, %groups_global : buffer, buffer, buffer, buffer, buffer + %groups_view = buffer.view %groups_na[%base] : buffer -> view<[%selected_count]xi32> + %qact_flat = buffer.view %qact_na[%base] : buffer -> view<1073741824xi8> + %activation_scale_view = buffer.view %scales_na[%base] : buffer -> view<33554432xf32> + %result0, %result1, %result2, %result3 = scf.if %valid_cohort -> (f32, f32, f32, f32) { + %acc0, %acc1, %acc2, %acc3 = scf.for %selected_index = [%c0 to %selected_count step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32, %iter3 = %zero_f32 : f32) -> (f32, f32, f32, f32) { + %group_i32 = view.load %groups_view[%selected_index] : view<[%selected_count]xi32> -> i32 + %group0 = index.cast %group_i32 : i32 to index + %group = index.assume %group0 [range(%group0, 0, 1023)] : index + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c8 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index = index.add %payload_index1, %row_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<8xi32> + %q2_payload = vector.load %weight_view[%c0] : view<8xi32> -> vector<8xi32> + %q2_00 = vector.extract %q2_payload[0] : vector<8xi32> -> i32 + %q2_01 = vector.extract %q2_payload[1] : vector<8xi32> -> i32 + %q2_10 = vector.extract %q2_payload[2] : vector<8xi32> -> i32 + %q2_11 = vector.extract %q2_payload[3] : vector<8xi32> -> i32 + %q2_20 = vector.extract %q2_payload[4] : vector<8xi32> -> i32 + %q2_21 = vector.extract %q2_payload[5] : vector<8xi32> -> i32 + %q2_30 = vector.extract %q2_payload[6] : vector<8xi32> -> i32 + %q2_31 = vector.extract %q2_payload[7] : vector<8xi32> -> i32 + %q2_packed0 = vector.from_elements %q2_00, %q2_01 : vector<2xi32> + %q2_packed1 = vector.from_elements %q2_10, %q2_11 : vector<2xi32> + %q2_packed2 = vector.from_elements %q2_20, %q2_21 : vector<2xi32> + %q2_packed3 = vector.from_elements %q2_30, %q2_31 : vector<2xi32> + %weight_words0 = func.call @ggml_expand_symmetric_i2_words(%q2_packed0) : (vector<2xi32>) -> (vector<4xi32>) + %weight_words1 = func.call @ggml_expand_symmetric_i2_words(%q2_packed1) : (vector<2xi32>) -> (vector<4xi32>) + %weight_words2 = func.call @ggml_expand_symmetric_i2_words(%q2_packed2) : (vector<2xi32>) -> (vector<4xi32>) + %weight_words3 = func.call @ggml_expand_symmetric_i2_words(%q2_packed3) : (vector<2xi32>) -> (vector<4xi32>) + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + %activation_payload_add = index.mul %group, %c16 : index + %activation_payload_index = index.assume %activation_payload_add [range(%activation_payload_add, 0, 1073741808)] : index + %activation_bytes = vector.load %qact_flat[%activation_payload_index] : view<1073741824xi8> -> vector<16xi8> + %activation_words = vector.bitcast %activation_bytes : vector<16xi8> to vector<4xi32> + %activation_scale_index = index.assume %group [range(%group, 0, 33554431)] : index + %activation_scale = view.load %activation_scale_view[%activation_scale_index] : view<33554432xf32> -> f32 + %combined_scale = scalar.mulf %weight_scale, %activation_scale : f32 + %next0 = func.call @ggml_symmetric_i2_dot_preloaded(%iter0, %weight_words0, %activation_words, %combined_scale) : (f32, vector<4xi32>, vector<4xi32>, f32) -> (f32) + %next1 = func.call @ggml_symmetric_i2_dot_preloaded(%iter1, %weight_words1, %activation_words, %combined_scale) : (f32, vector<4xi32>, vector<4xi32>, f32) -> (f32) + %next2 = func.call @ggml_symmetric_i2_dot_preloaded(%iter2, %weight_words2, %activation_words, %combined_scale) : (f32, vector<4xi32>, vector<4xi32>, f32) -> (f32) + %next3 = func.call @ggml_symmetric_i2_dot_preloaded(%iter3, %weight_words3, %activation_words, %combined_scale) : (f32, vector<4xi32>, vector<4xi32>, f32) -> (f32) + scf.yield %next0, %next1, %next2, %next3 : f32, f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2, %acc3 : f32, f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32, f32 + } + scf.if %valid_cohort { + %output_view = buffer.view %output_na[%base] : buffer -> view<[%output_size]xf32> + %results = vector.from_elements %result0, %result1, %result2, %result3 : vector<4xf32> + vector.store %results, %output_view[%row] : vector<4xf32>, view<[%output_size]xf32> + } + kernel.return +} + +kernel.def export("ggml_top_k8_f32_partitions_register") @ggml_top_k8_f32_partitions_register(%element_count: index) { + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c128, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%element_count: index, %values: buffer, %partial_values: buffer, %partial_ids: buffer) { + %count = index.assume %element_count [range(%element_count, 64, 262144)] : index + %partition = kernel.workgroup.id : index + %tid = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %candidate_count = index.constant 8 : index + %c8 = index.constant 8 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %negative_large = scalar.constant -3.4028234663852886e+38 : f32 + %largest_i32 = scalar.constant 2147483647 : i32 + %negative_values = vector.splat %negative_large : vector<8xf32> + %invalid_ids = vector.splat %largest_i32 : vector<8xi32> + %values_noalias, %partial_values_noalias, %partial_ids_noalias = buffer.assume.noalias %values, %partial_values, %partial_ids : buffer, buffer, buffer + %values_view = buffer.view %values_noalias[%base] : buffer -> view<[%count]xf32> + %partial_values_view = buffer.view %partial_values_noalias[%base] : buffer -> view<1024xf32> + %partial_ids_view = buffer.view %partial_ids_noalias[%base] : buffer -> view<1024xi32> + %rounded_count = index.add %count, %c127 : index + %partition_span = index.div %rounded_count, %c128 : index + %partition_begin = index.mul %partition, %partition_span : index + %partition_end_unbounded = index.add %partition_begin, %partition_span : index + %end_exceeds_count = index.cmp ugt, %partition_end_unbounded, %count : index + %partition_end = scf.select %end_exceeds_count, %count, %partition_end_unbounded : index + %thread_begin = index.add %partition_begin, %tid : index + %partial_base = index.mul %partition, %candidate_count : index + %loaded_values, %loaded_ids = scf.for %slot = [%c0 to %c8 step %c1](%local_values = %negative_values : vector<8xf32>, %local_ids = %invalid_ids : vector<8xi32>) -> (vector<8xf32>, vector<8xi32>) unroll { + %slot_stride = index.mul %slot, %c256 : index + %candidate_index = index.add %thread_begin, %slot_stride : index + %valid = index.cmp ult, %candidate_index, %partition_end : index + %candidate_value, %candidate_id = scf.if %valid -> (f32, i32) { + %bounded_index = index.assume %candidate_index [lt(%candidate_index, %count)] : index + %value = view.load %values_view[%bounded_index] : view<[%count]xf32> -> f32 + %id = index.cast %bounded_index : index to i32 + scf.yield %value, %id : f32, i32 + } else { + scf.yield %negative_large, %largest_i32 : f32, i32 + } + %next_values = vector.insert %candidate_value into %local_values[%slot] : f32, vector<8xf32> + %next_ids = vector.insert %candidate_id into %local_ids[%slot] : i32, vector<8xi32> + scf.yield %next_values, %next_ids : vector<8xf32>, vector<8xi32> + } + %remaining_final = scf.for %rank = [%c0 to %candidate_count step %c1](%remaining = %loaded_values : vector<8xf32>) -> (vector<8xf32>) unroll { + %local_value, %local_id = scf.for %slot = [%c0 to %c8 step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { + %candidate_value = vector.extract %remaining[%slot] : vector<8xf32> -> f32 + %candidate_id = vector.extract %loaded_ids[%slot] : vector<8xi32> -> i32 + %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 + %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 + %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 + %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 + %is_better = scalar.ori %is_greater, %is_lower_tie : i1 + %next_value = scf.select %is_better, %candidate_value, %best_value : f32 + %next_id = scf.select %is_better, %candidate_id, %best_id : i32 + scf.yield %next_value, %next_id : f32, i32 + } + %winner_value = kernel.workgroup.reduce %local_value : f32 + %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 + %winner_id_candidate = scf.select %matches_winner, %local_id, %largest_i32 : i32 + %winner_id = kernel.workgroup.reduce %winner_id_candidate : i32 + %is_tid_zero = index.cmp eq, %tid, %c0 : index + scf.if %is_tid_zero { + %partial_index = index.add %partial_base, %rank : index + view.store %winner_value, %partial_values_view[%partial_index] : f32, view<1024xf32> + view.store %winner_id, %partial_ids_view[%partial_index] : i32, view<1024xi32> + } + %next_remaining = scf.for %slot = [%c0 to %c8 step %c1](%current = %remaining : vector<8xf32>) -> (vector<8xf32>) unroll { + %candidate_id = vector.extract %loaded_ids[%slot] : vector<8xi32> -> i32 + %is_winner = scalar.cmpi eq, %candidate_id, %winner_id : i32 + %old_value = vector.extract %current[%slot] : vector<8xf32> -> f32 + %next_value = scf.select %is_winner, %negative_large, %old_value : f32 + %updated = vector.insert %next_value into %current[%slot] : f32, vector<8xf32> + scf.yield %updated : vector<8xf32> + } + scf.yield %next_remaining : vector<8xf32> + } + kernel.return +} + +kernel.def export("ggml_top_k128_f32_reduce_gather_register") @ggml_top_k128_f32_reduce_gather_register(%element_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%element_count: index, %partial_values: buffer, %partial_ids: buffer, %candidate_output: buffer, %value_output: buffer) { + %count = index.assume %element_count [range(%element_count, 64, 262144)] : index + %tid = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %candidate_count = index.constant 128 : index + %c256 = index.constant 256 : index + %c1024 = index.constant 1024 : index + %base = index.constant 0 : offset + %negative_large = scalar.constant -3.4028234663852886e+38 : f32 + %largest_i32 = scalar.constant 2147483647 : i32 + %negative_values = vector.splat %negative_large : vector<4xf32> + %invalid_ids = vector.splat %largest_i32 : vector<4xi32> + %partial_values_noalias, %partial_ids_noalias, %candidate_output_noalias, %value_output_noalias = buffer.assume.noalias %partial_values, %partial_ids, %candidate_output, %value_output : buffer, buffer, buffer, buffer + %partial_values_view = buffer.view %partial_values_noalias[%base] : buffer -> view<1024xf32> + %partial_ids_view = buffer.view %partial_ids_noalias[%base] : buffer -> view<1024xi32> + %candidate_output_view = buffer.view %candidate_output_noalias[%base] : buffer -> view<128xi32> + %value_output_view = buffer.view %value_output_noalias[%base] : buffer -> view<128xf32> + %loaded_values, %loaded_ids = scf.for %slot = [%c0 to %c4 step %c1](%values = %negative_values : vector<4xf32>, %ids = %invalid_ids : vector<4xi32>) -> (vector<4xf32>, vector<4xi32>) unroll { + %slot_stride = index.mul %slot, %c256 : index + %partial_index = index.add %tid, %slot_stride : index + %valid = index.cmp ult, %partial_index, %c1024 : index + %candidate_value, %candidate_id = scf.if %valid -> (f32, i32) { + %value = view.load %partial_values_view[%partial_index] : view<1024xf32> -> f32 + %id = view.load %partial_ids_view[%partial_index] : view<1024xi32> -> i32 + scf.yield %value, %id : f32, i32 + } else { + scf.yield %negative_large, %largest_i32 : f32, i32 + } + %next_values = vector.insert %candidate_value into %values[%slot] : f32, vector<4xf32> + %next_ids = vector.insert %candidate_id into %ids[%slot] : i32, vector<4xi32> + scf.yield %next_values, %next_ids : vector<4xf32>, vector<4xi32> + } + %remaining_final = scf.for %rank = [%c0 to %candidate_count step %c1](%remaining = %loaded_values : vector<4xf32>) -> (vector<4xf32>) { + %local_value, %local_id = scf.for %slot = [%c0 to %c4 step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { + %candidate_value = vector.extract %remaining[%slot] : vector<4xf32> -> f32 + %candidate_id = vector.extract %loaded_ids[%slot] : vector<4xi32> -> i32 + %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 + %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 + %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 + %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 + %is_better = scalar.ori %is_greater, %is_lower_tie : i1 + %next_value = scf.select %is_better, %candidate_value, %best_value : f32 + %next_id = scf.select %is_better, %candidate_id, %best_id : i32 + scf.yield %next_value, %next_id : f32, i32 + } + %winner_value = kernel.workgroup.reduce %local_value : f32 + %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 + %winner_id_candidate = scf.select %matches_winner, %local_id, %largest_i32 : i32 + %winner_id = kernel.workgroup.reduce %winner_id_candidate : i32 + %is_tid_zero = index.cmp eq, %tid, %c0 : index + scf.if %is_tid_zero { + view.store %winner_id, %candidate_output_view[%rank] : i32, view<128xi32> + view.store %winner_value, %value_output_view[%rank] : f32, view<128xf32> + } + %next_remaining = scf.for %slot = [%c0 to %c4 step %c1](%values = %remaining : vector<4xf32>) -> (vector<4xf32>) unroll { + %candidate_id = vector.extract %loaded_ids[%slot] : vector<4xi32> -> i32 + %is_winner = scalar.cmpi eq, %candidate_id, %winner_id : i32 + %old_value = vector.extract %values[%slot] : vector<4xf32> -> f32 + %next_value = scf.select %is_winner, %negative_large, %old_value : f32 + %updated = vector.insert %next_value into %values[%slot] : f32, vector<4xf32> + scf.yield %updated : vector<4xf32> + } + scf.yield %next_remaining : vector<4xf32> + } + kernel.return +} + +kernel.def export("ggml_fill_negative_f32") @ggml_fill_negative_f32(%element_count: index) { + %c1 = index.constant 1 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %rounded_count = index.add %element_count, %c255 : index + %workgroups = index.div %rounded_count, %c256 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%element_count: index, %output: buffer) { + %count = index.assume %element_count [range(%element_count, 1, 262144)] : index + %workgroup = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %negative_large = scalar.constant -3.4028234663852886e+38 : f32 + %workgroup_base = index.mul %workgroup, %c256 : index + %index0 = index.add %workgroup_base, %tid : index + %index = index.assume %index0 [range(%index0, 0, 262399)] : index + %valid = index.cmp ult, %index, %count : index + %output_view = buffer.view %output[%base] : buffer -> view<[%count]xf32> + scf.if %valid { + %bounded_index = index.assume %index [lt(%index, %count)] : index + view.store %negative_large, %output_view[%bounded_index] : f32, view<[%count]xf32> + } + kernel.return +} + +template.def<@ggml.mul_mat_q6_k_packed.selected_refine.body> device @ggml_mul_mat_q6_k_packed_selected_refine_body(%token_count: index, %input_size0: index, %output_size0: index, %scale_row_layout: index, %candidate_count0: index, %input: buffer, %weight: buffer, %candidates: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index + %weight_offset_index = config.get @ggml.mul_mat_q6_k_packed.weight_offset : index + %weight_offset = index.cast %weight_offset_index : index to offset + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 16)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %bounded_candidate_count = index.assume %candidate_count0 [range(%candidate_count0, 1, 64)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c72 = index.constant 72 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %c128_bytes = index.constant 128 : offset + %c1152_bytes = index.constant 1152 : offset + %c13440_bytes = index.constant 13440 : offset + %weight_stage_bytes = index.constant 9216 : offset + %activation_stage_bytes = index.constant 2304 : offset + %result_stage_bytes = index.constant 2048 : offset + %zero_accumulator = vector.constant 0.0 : vector<16xf16> + %empty_words = vector.constant 0 : vector<4xi32> + %zero_scale_bytes = vector.constant 0 : vector<8xi8> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high0_mask = vector.constant 50529027 : vector<4xi32> + %high1_mask = vector.constant 202116108 : vector<4xi32> + %high2_mask = vector.constant 808464432 : vector<4xi32> + %high3_mask = vector.constant -1061109568 : vector<4xi32> + %shift2 = vector.constant 2 : vector<4xi32> + %shift4 = vector.constant 4 : vector<4xi32> + %negative32_f32 = vector.constant -32.0 : vector<16xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %quant_block_count = index.div %bounded_input_size, %c256 : index + %input_noalias, %weight_noalias, %candidates_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %candidates, %output : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%base] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %candidates_view = buffer.view %candidates_noalias[%base] : buffer -> view<64xi32> + %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %weight_stage_view = buffer.view %weight_stage[%base] : buffer -> view<64x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%base] : buffer -> view<16x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%base] : buffer -> view<64x16xf16, %activation_fragment_layout> + %weight_half = index.rem %workitem, %c2 : index + %weight_local_row0 = index.div %workitem, %c2 : index + %weight_local_row = index.assume %weight_local_row0 [range(%weight_local_row0, 0, 63)] : index + %weight_k = index.mul %weight_half, %c16 : index + %weight_k_next = index.add %weight_k, %c32 : index + %valid_candidate = index.cmp ult, %weight_local_row, %bounded_candidate_count : index + %channel0 = scf.if %valid_candidate -> (index) { + %candidate_i32 = view.load %candidates_view[%weight_local_row] : view<64xi32> -> i32 + %candidate_index0 = index.cast %candidate_i32 : i32 to index + %candidate_index = index.assume %candidate_index0 [range(%candidate_index0, 0, 262143)] : index + scf.yield %candidate_index : index + } else { + scf.yield %c0 : index + } + %channel = index.assume %channel0 [range(%channel0, 0, 262143)] : index + %channel_in_range = index.cmp ult, %channel, %bounded_output_size : index + %valid_channel = scalar.andi %valid_candidate, %channel_in_range : i1 + %selected_channel_tile = index.div %channel, %c64 : index + %selected_weight_local_row = index.rem %channel, %c64 : index + %use_scale_row_layout = index.cmp eq, %scale_row_layout, %c1 : index + %activation_row0 = index.div %workitem, %c8 : index + %activation_row = index.assume %activation_row0 [range(%activation_row0, 0, 15)] : index + %activation_packet = index.rem %workitem, %c8 : index + %activation_k = index.mul %activation_packet, %c4 : index + %activation_k_next = index.add %activation_k, %c32 : index + %valid_token = index.cmp ult, %activation_row, %bounded_token_count : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %init = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> + %result = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc = %init : vector<16xf16>) -> (vector<16xf16>) { + %row_group_block0 = index.mul %selected_channel_tile, %quant_block_count : index + %row_group_block = index.add %row_group_block0, %quant_block : index + %tile_byte_relative = index.scale %row_group_block, %c13440_bytes : index, offset -> offset + %tile_byte_base = index.add %weight_offset, %tile_byte_relative : offset + %scale_byte_base = index.add %tile_byte_base, %c128_bytes : offset + %raw_byte_base = index.add %tile_byte_base, %c1152_bytes : offset + %d_view = buffer.view %weight_noalias[%tile_byte_base] : buffer -> view<64xf16> + %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<8x64x2xi8> + %scale_row_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<64x2x8xi8> + %raw_view = buffer.view %weight_noalias[%raw_byte_base] : buffer -> view<2x3x64x32xi8> + %d = scf.if %valid_channel -> (f32) { + %d_f16 = view.load %d_view[%selected_weight_local_row] : view<64xf16> -> f16 + %d_f32 = scalar.extf %d_f16 : f16 to f32 + scf.yield %d_f32 : f32 + } else { + %zero_f32 = scalar.constant 0.0 : f32 + scf.yield %zero_f32 : f32 + } + %load_scale_row = scalar.andi %valid_channel, %use_scale_row_layout : i1 + %scale_bytes = scf.if %load_scale_row -> (vector<8xi8>) { + %loaded_scale_bytes = vector.load %scale_row_view[%selected_weight_local_row, %weight_half, %c0] : view<64x2x8xi8> -> vector<8xi8> + scf.yield %loaded_scale_bytes : vector<8xi8> + } else { + scf.yield %zero_scale_bytes : vector<8xi8> + } + %block_result, %final_ql0_words, %final_ql1_words, %final_qh_words = scf.for %quant_group = [%c0 to %c8 step %c2](%acc = %block_acc : vector<16xf16>, %cached_ql0_words = %empty_words : vector<4xi32>, %cached_ql1_words = %empty_words : vector<4xi32>, %cached_qh_words = %empty_words : vector<4xi32>) -> (vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { + %quant_group_next = index.add %quant_group, %c1 : index + %half128 = index.div %quant_group, %c4 : index + %group_in_half = index.rem %quant_group, %c4 : index + %is_low_pair = index.cmp eq, %group_in_half, %c0 : index + %is_first_group = index.cmp eq, %quant_group, %c0 : index + %prefetch_second_half = index.cmp eq, %quant_group, %c2 : index + %weight_values0, %weight_values1, %next_ql0_words, %next_ql1_words, %next_qh_words = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_words, %ql1_words, %qh_words = scf.if %is_first_group -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %ql0_bytes = vector.load %raw_view[%half128, %c0, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %ql1_bytes = vector.load %raw_view[%half128, %c1, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %qh_bytes = vector.load %raw_view[%half128, %c2, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %loaded_ql0_words = vector.bitcast %ql0_bytes : vector<16xi8> to vector<4xi32> + %loaded_ql1_words = vector.bitcast %ql1_bytes : vector<16xi8> to vector<4xi32> + %loaded_qh_words = vector.bitcast %qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %loaded_ql0_words, %loaded_ql1_words, %loaded_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + %codes0_words, %codes1_words = scf.if %is_low_pair -> (vector<4xi32>, vector<4xi32>) { + %ql0_low = vector.andi %ql0_words, %nibble_mask : vector<4xi32> + %ql1_low = vector.andi %ql1_words, %nibble_mask : vector<4xi32> + %qh0_bits0 = vector.andi %qh_words, %high0_mask : vector<4xi32> + %qh1_bits0 = vector.andi %qh_words, %high1_mask : vector<4xi32> + %qh0_bits = vector.shli %qh0_bits0, %shift4 : vector<4xi32> + %qh1_bits = vector.shli %qh1_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_low, %qh0_bits : vector<4xi32> + %codes1 = vector.ori %ql1_low, %qh1_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } else { + %ql0_high0 = vector.shrui %ql0_words, %shift4 : vector<4xi32> + %ql1_high0 = vector.shrui %ql1_words, %shift4 : vector<4xi32> + %ql0_high = vector.andi %ql0_high0, %nibble_mask : vector<4xi32> + %ql1_high = vector.andi %ql1_high0, %nibble_mask : vector<4xi32> + %qh2_bits = vector.andi %qh_words, %high2_mask : vector<4xi32> + %qh3_bits0 = vector.andi %qh_words, %high3_mask : vector<4xi32> + %qh3_bits = vector.shrui %qh3_bits0, %shift2 : vector<4xi32> + %codes0 = vector.ori %ql0_high, %qh2_bits : vector<4xi32> + %codes1 = vector.ori %ql1_high, %qh3_bits : vector<4xi32> + scf.yield %codes0, %codes1 : vector<4xi32>, vector<4xi32> + } + %codes0_u8 = vector.bitcast %codes0_words : vector<4xi32> to vector<16xi8> + %codes1_u8 = vector.bitcast %codes1_words : vector<4xi32> to vector<16xi8> + %codes0_unsigned_f32 = vector.uitofp %codes0_u8 : vector<16xi8> to vector<16xf32> + %codes1_unsigned_f32 = vector.uitofp %codes1_u8 : vector<16xi8> to vector<16xf32> + %codes0_f32 = vector.addf %codes0_unsigned_f32, %negative32_f32 : vector<16xf32> + %codes1_f32 = vector.addf %codes1_unsigned_f32, %negative32_f32 : vector<16xf32> + %scale0_i8, %scale1_i8 = scf.if %use_scale_row_layout -> (i8, i8) { + %row_scale0_i8 = vector.extract %scale_bytes[%quant_group] : vector<8xi8> -> i8 + %row_scale1_i8 = vector.extract %scale_bytes[%quant_group_next] : vector<8xi8> -> i8 + scf.yield %row_scale0_i8, %row_scale1_i8 : i8, i8 + } else { + %group_scale0_i8 = view.load %scale_view[%quant_group, %selected_weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + %group_scale1_i8 = view.load %scale_view[%quant_group_next, %selected_weight_local_row, %weight_half] : view<8x64x2xi8> -> i8 + scf.yield %group_scale0_i8, %group_scale1_i8 : i8, i8 + } + %scale0 = scalar.sitofp %scale0_i8 : i8 to f32 + %scale1 = scalar.sitofp %scale1_i8 : i8 to f32 + %combined0 = scalar.mulf %scale0, %d : f32 + %combined1 = scalar.mulf %scale1, %d : f32 + %combined0_vector = vector.splat %combined0 : vector<16xf32> + %combined1_vector = vector.splat %combined1 : vector<16xf32> + %values0_f32 = vector.mulf %codes0_f32, %combined0_vector : vector<16xf32> + %values1_f32 = vector.mulf %codes1_f32, %combined1_vector : vector<16xf32> + %values0 = vector.fptrunc %values0_f32 : vector<16xf32> to vector<16xf16> + %values1 = vector.fptrunc %values1_f32 : vector<16xf32> to vector<16xf16> + %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words = scf.if %prefetch_second_half -> (vector<4xi32>, vector<4xi32>, vector<4xi32>) { + %next_ql0_bytes = vector.load %raw_view[%c1, %c0, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql1_bytes = vector.load %raw_view[%c1, %c1, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_qh_bytes = vector.load %raw_view[%c1, %c2, %selected_weight_local_row, %weight_k] : view<2x3x64x32xi8> -> vector<16xi8> + %next_ql0 = vector.bitcast %next_ql0_bytes : vector<16xi8> to vector<4xi32> + %next_ql1 = vector.bitcast %next_ql1_bytes : vector<16xi8> to vector<4xi32> + %next_qh = vector.bitcast %next_qh_bytes : vector<16xi8> to vector<4xi32> + scf.yield %next_ql0, %next_ql1, %next_qh : vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + scf.yield %ql0_words, %ql1_words, %qh_words : vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %values0, %values1, %prefetched_ql0_words, %prefetched_ql1_words, %prefetched_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } else { + %zeros0 = vector.constant 0.0 : vector<16xf16> + %zeros1 = vector.constant 0.0 : vector<16xf16> + scf.yield %zeros0, %zeros1, %cached_ql0_words, %cached_ql1_words, %cached_qh_words : vector<16xf16>, vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %weight_k] : vector<16xf16>, view<64x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k_next] : vector<16xf16>, view<64x72xf16> + %activation_values0, %activation_values1 = scf.if %valid_token -> (vector<4xf16>, vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %activation_row, %bounded_token_count [lt(%activation_row, %bounded_token_count)] : index, index + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %input_k0 = index.add %block_k_base, %group_k_add : index + %input_k = index.add %input_k0, %activation_k : index + %input_k_next = index.add %input_k0, %activation_k_next : index + %loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> + %narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16> + %narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16> + scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16> + } else { + %zeros0 = vector.constant 0.0 : vector<4xf16> + %zeros1 = vector.constant 0.0 : vector<4xf16> + scf.yield %zeros0, %zeros1 : vector<4xf16>, vector<4xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_row, %activation_k] : vector<4xf16>, view<16x72xf16> + vector.store %activation_values1, %activation_stage_physical_view[%activation_row, %activation_k_next] : vector<4xf16>, view<16x72xf16> + kernel.barrier scope(workgroup) ordering(acq_rel) + %next = scf.for %k_half = [%c0 to %c64 step %c16](%half_acc = %acc : vector<16xf16>) -> (vector<16xf16>) unroll { + %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x72xf16> -> vector<16xf16> + %rhs = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<64x16xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next = vector.mma %lhs, %rhs, %half_acc : vector<16xf16>, vector<16xf16>, vector<16xf16> + scf.yield %half_next : vector<16xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next, %next_ql0_words, %next_ql1_words, %next_qh_words : vector<16xf16>, vector<4xi32>, vector<4xi32>, vector<4xi32> + } + scf.yield %block_result : vector<16xf16> + } + template.apply<@ggml.mul_mat_q6_k_packed.publish_f16_wave32>(%result, %c0, %c0, %bounded_token_count, %c64, %c0, %subgroup_channel_add, %result_stage, %output_noalias, %output_accumulation) : (vector<16xf16>, index, index, index, index, index, index, buffer, buffer, index) + template.return +} + +kernel.def export("ggml_mul_mat_q6_k_packed_selected_refine_token1") @ggml_mul_mat_q6_k_packed_selected_refine_token1(%token_count: index, %candidate_count: index) { + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %candidate_count: index, %input: buffer, %weight: buffer, %candidates: buffer, %exact_output: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_q6_k_packed.input_size : index + %output_size0 = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c64 = index.constant 64 : index + %base = index.constant 0 : offset + %scale_row_layout = index.constant 1 : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %bounded_candidate_count = index.assume %candidate_count [range(%candidate_count, 1, 64)] : index + template.apply<@ggml.mul_mat_q6_k_packed.selected_refine.body>(%bounded_token_count, %input_size, %output_size, %scale_row_layout, %bounded_candidate_count, %input, %weight, %candidates, %exact_output) : (index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.barrier scope(workgroup) ordering(acq_rel) + %tid0 = kernel.workitem.id : index + %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index + %publishes = index.cmp ult, %tid, %bounded_candidate_count : index + %candidates_noalias, %exact_noalias, %output_noalias = buffer.assume.noalias %candidates, %exact_output, %output : buffer, buffer, buffer + %candidate_view = buffer.view %candidates_noalias[%base] : buffer -> view<64xi32> + %exact_view = buffer.view %exact_noalias[%base] : buffer -> view<64xf32> + scf.if %publishes { + %candidate_i32 = view.load %candidate_view[%tid] : view<64xi32> -> i32 + %candidate0 = index.cast %candidate_i32 : i32 to index + %candidate1 = index.assume %candidate0 [range(%candidate0, 0, 262143)] : index + %candidate, %bounded_output_size = index.assume %candidate1, %output_size [lt(%candidate1, %output_size)] : index, index + %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%bounded_output_size]xf32> + %value = view.load %exact_view[%tid] : view<64xf32> -> f32 + view.store %value, %output_view[%candidate] : f32, view<[%bounded_output_size]xf32> + } + kernel.return +} + +template.decl @ggml.mul_mat_q6_k_packed.publish_f32_wave32(%result: vector<8xf32>, %token_offset: index, %token_tile_base: index, %token_count: index, %output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %output: buffer, %output_accumulation: index) + +template.def<@ggml.mul_mat_q6_k_packed.publish_f32_wave32> device @ggml_mul_mat_q6_k_packed_publish_f32_wave32(%result: vector<8xf32>, %token_offset: index, %token_tile_base: index, %token_count: index, %output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %output: buffer, %output_accumulation: index) { + %base = index.constant 0 : offset + %c1 = index.constant 1 : index + %c16 = index.constant 16 : index + %token = index.add %token_tile_base, %token_offset : index + %channel = index.add %channel_tile_base, %subgroup_channel_add : index + %token_end = index.add %token, %c16 : index + %channel_end = index.add %channel, %c16 : index + %valid_token = index.cmp ule, %token_end, %token_count : index + %valid_channel = index.cmp ule, %channel_end, %output_size : index + %valid = scalar.andi %valid_token, %valid_channel : i1 + scf.if %valid { + %layout = encoding.layout.strided [1, %output_size] : encoding + %dst = buffer.view %output[%base] : buffer -> view<[%output_size]x[%token_count]xf32, %layout> + %accumulates = index.cmp eq, %output_accumulation, %c1 : index + %published = scf.if %accumulates -> (vector<8xf32>) { + %old = vector.fragment.load %dst[%channel, %token] shape [%c16, %c16] : view<[%output_size]x[%token_count]xf32, %layout> -> vector<8xf32> + %sum = vector.addf %old, %result : vector<8xf32> + scf.yield %sum : vector<8xf32> + } else { + scf.yield %result : vector<8xf32> + } + vector.fragment.store %published, %dst[%channel, %token] shape [%c16, %c16] : vector<8xf32>, view<[%output_size]x[%token_count]xf32, %layout> + } else { + // Native wave32 F32 result: lane % 16 selects the token; the high + // lane bit selects even/odd channels in the eight-value fragment. + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %lane = kernel.subgroup.lane.id : index + %lane_token = index.rem %lane, %c16 : index + %lane_channel = index.div %lane, %c16 : index + %partial_token = index.add %token, %lane_token : index + %partial_channel_base = index.add %channel, %lane_channel : index + %token_valid = index.cmp ult, %partial_token, %token_count : index + %accumulates = index.cmp eq, %output_accumulation, %c1 : index + scf.if %token_valid { + %bounded_token, %output_tokens = index.assume %partial_token, %token_count [lt(%partial_token, %token_count)] : index, index + %dst = buffer.view %output[%base] : buffer -> view<[%output_tokens]x[%output_size]xf32> + scf.for %element = [%c0 to %c8 step %c1] unroll { + %channel_step = index.mul %element, %c2 : index + %partial_channel = index.add %partial_channel_base, %channel_step : index + %channel_valid = index.cmp ult, %partial_channel, %output_size : index + scf.if %channel_valid { + %bounded_channel, %output_channels = index.assume %partial_channel, %output_size [lt(%partial_channel, %output_size)] : index, index + %value = vector.extract %result[%element] : vector<8xf32> -> f32 + %published = scf.if %accumulates -> (f32) { + %old = view.load %dst[%bounded_token, %bounded_channel] : view<[%output_tokens]x[%output_size]xf32> -> f32 + %sum = scalar.addf %value, %old : f32 + scf.yield %sum : f32 + } else { + scf.yield %value : f32 + } + view.store %published, %dst[%bounded_token, %bounded_channel] : f32, view<[%output_tokens]x[%output_size]xf32> + } + } + } + } + template.return +} + +template.def<@ggml.mul_mat_q6_k_packed.prefill_wave32.body> device @ggml_mul_mat_q6_k_packed_prefill_wave32_body(%input_is_f16: i1, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %output: buffer) { + %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index + %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, 256)] : 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, 7)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c72 = index.constant 72 : index + %c48 = index.constant 48 : index + %c64 = index.constant 64 : index + %c80 = index.constant 80 : index + %c96 = index.constant 96 : index + %c112 = index.constant 112 : index + %c128 = index.constant 128 : index + %c210 = index.constant 210 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %weight_stage_bytes = index.constant 18432 : offset + %activation_stage_bytes = index.constant 9216 : offset + %c0_f16x4 = vector.constant 0.0 : vector<16xf16> + %c0_f16x8 = vector.constant 0.0 : vector<8xf16> + %zero_accumulator = vector.constant 0.0 : vector<8xf32> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %quant_block_count = index.div %bounded_input_size, %c256 : index + %weight_row_bytes = index.mul %quant_block_count, %c210 : index + %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, 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 + %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<128x72xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x72xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c72] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<64x64xf16, %activation_fragment_layout> + %channel_tile_base = index.mul %channel_tile, %c128 : index + %token_tile_base = index.mul %token_tile, %c128 : index + %second_token_tile_base = index.add %token_tile_base, %c64 : index + %load_half = index.rem %workitem, %c2 : index + %load_packet = index.mul %load_half, %c4 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c2 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 127)] : index + %activation_load_packet = index.rem %workitem, %c4 : index + %activation_load_k = index.mul %activation_load_packet, %c8 : index + %activation_load_row0 = index.div %workitem, %c4 : index + %activation_load_row = index.assume %activation_load_row0 [range(%activation_load_row0, 0, 63)] : index + %subgroup_channel_add = index.mul %subgroup, %c16 : index + %rhs_column0 = index.add %c0, %c0 : index + %rhs_column1 = index.add %c16, %c0 : index + %rhs_column2 = index.add %c32, %c0 : index + %rhs_column3 = index.add %c48, %c0 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init02 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init03 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init12 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %init13 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf32> + %result00, %result01, %result02, %result03, %result10, %result11, %result12, %result13 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf32>, %block_acc01 = %init01 : vector<8xf32>, %block_acc02 = %init02 : vector<8xf32>, %block_acc03 = %init03 : vector<8xf32>, %block_acc10 = %init10 : vector<8xf32>, %block_acc11 = %init11 : vector<8xf32>, %block_acc12 = %init12 : vector<8xf32>, %block_acc13 = %init13 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %block_result00, %block_result01, %block_result02, %block_result03, %block_result10, %block_result11, %block_result12, %block_result13 = scf.for %quant_group = [%c0 to %c8 step %c2](%acc00 = %block_acc00 : vector<8xf32>, %acc01 = %block_acc01 : vector<8xf32>, %acc02 = %block_acc02 : vector<8xf32>, %acc03 = %block_acc03 : vector<8xf32>, %acc10 = %block_acc10 : vector<8xf32>, %acc11 = %block_acc11 : vector<8xf32>, %acc12 = %block_acc12 : vector<8xf32>, %acc13 = %block_acc13 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %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 + %weight_local_row = index.assume %load_row [range(%load_row, 0, 127)] : index + %channel = index.add %channel_tile_base, %weight_local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { + %channel_byte_index = index.mul %channel, %weight_row_bytes : index + %row_byte_base = index.cast %channel_byte_index : index to offset + %half_packet = scalar.constant false : i1 + %decoded0, %decoded1 = func.call @ggml_q6k_f16_pair(%half_packet, %weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> + } else { + scf.yield %c0_f16x4, %c0_f16x4 : vector<16xf16>, vector<16xf16> + } + %weight_k1 = index.add %load_k, %c32 : index + vector.store %weight_values0, %weight_stage_view[%weight_local_row, %load_k] : vector<16xf16>, view<128x72xf16> + vector.store %weight_values1, %weight_stage_view[%weight_local_row, %weight_k1] : vector<16xf16>, view<128x72xf16> + scf.for %stage_half0 = [%c0 to %c2 step %c1] unroll { + %stage_k_base0 = index.mul %stage_half0, %c32 : index + %token0 = index.add %token_tile_base, %activation_load_row : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %activation_k0 = index.add %stage_k_base0, %activation_load_k : index + %activation_values0 = scf.if %valid_token0 -> (vector<8xf16>) { + %bounded_token0, %input_token_count0 = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %input_k0 = index.add %k_origin, %activation_k0 : index + %loaded0 = scf.if %input_is_f16 -> (vector<8xf16>) { + %narrow0 = vector.load %input_f16_view[%bounded_token0, %input_k0] : view<[%bounded_token_count]x[%bounded_input_size]xf16> -> vector<8xf16> + scf.yield %narrow0 : vector<8xf16> + } else { + %wide0 = vector.load %input_f32_view[%bounded_token0, %input_k0] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<8xf32> + %narrow0 = vector.fptrunc %wide0 : vector<8xf32> to vector<8xf16> + scf.yield %narrow0 : vector<8xf16> + } + scf.yield %loaded0 : vector<8xf16> + } else { + scf.yield %c0_f16x8 : vector<8xf16> + } + vector.store %activation_values0, %activation_stage_physical_view[%activation_load_row, %activation_k0] : vector<8xf16>, view<64x72xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next00, %next01, %next02, %next03 = scf.for %k_half0 = [%c0 to %c64 step %c16](%half_acc00 = %acc00 : vector<8xf32>, %half_acc01 = %acc01 : vector<8xf32>, %half_acc02 = %acc02 : vector<8xf32>, %half_acc03 = %acc03 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) unroll { + %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half0] shape [%m, %k] : view<128x72xf16> -> vector<16xf16> + %rhs00 = vector.fragment.load %activation_fragment_view[%k_half0, %c0] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs01 = vector.fragment.load %activation_fragment_view[%k_half0, %c16] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next00 = vector.mma %lhs0, %rhs00, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next01 = vector.mma %lhs0, %rhs01, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.schedule.fence + %rhs02 = vector.fragment.load %activation_fragment_view[%k_half0, %c32] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs03 = vector.fragment.load %activation_fragment_view[%k_half0, %c48] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next02 = vector.mma %lhs0, %rhs02, %half_acc02 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next03 = vector.mma %lhs0, %rhs03, %half_acc03 : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.schedule.fence + scf.yield %half_next00, %half_next01, %half_next02, %half_next03 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %stage_half1 = [%c0 to %c2 step %c1] unroll { + %stage_k_base1 = index.mul %stage_half1, %c32 : index + %token1 = index.add %second_token_tile_base, %activation_load_row : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %activation_k1 = index.add %stage_k_base1, %activation_load_k : index + %activation_values1 = scf.if %valid_token1 -> (vector<8xf16>) { + %bounded_token1, %input_token_count1 = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index + %input_k1 = index.add %k_origin, %activation_k1 : index + %loaded1 = scf.if %input_is_f16 -> (vector<8xf16>) { + %narrow1 = vector.load %input_f16_view[%bounded_token1, %input_k1] : view<[%bounded_token_count]x[%bounded_input_size]xf16> -> vector<8xf16> + scf.yield %narrow1 : vector<8xf16> + } else { + %wide1 = vector.load %input_f32_view[%bounded_token1, %input_k1] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<8xf32> + %narrow1 = vector.fptrunc %wide1 : vector<8xf32> to vector<8xf16> + scf.yield %narrow1 : vector<8xf16> + } + scf.yield %loaded1 : vector<8xf16> + } else { + scf.yield %c0_f16x8 : vector<8xf16> + } + vector.store %activation_values1, %activation_stage_physical_view[%activation_load_row, %activation_k1] : vector<8xf16>, view<64x72xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %next10, %next11, %next12, %next13 = scf.for %k_half1 = [%c0 to %c64 step %c16](%half_acc10 = %acc10 : vector<8xf32>, %half_acc11 = %acc11 : vector<8xf32>, %half_acc12 = %acc12 : vector<8xf32>, %half_acc13 = %acc13 : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) unroll { + %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half1] shape [%m, %k] : view<128x72xf16> -> vector<16xf16> + %rhs10 = vector.fragment.load %activation_fragment_view[%k_half1, %c0] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs11 = vector.fragment.load %activation_fragment_view[%k_half1, %c16] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next10 = vector.mma %lhs1, %rhs10, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next11 = vector.mma %lhs1, %rhs11, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.schedule.fence + %rhs12 = vector.fragment.load %activation_fragment_view[%k_half1, %c32] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs13 = vector.fragment.load %activation_fragment_view[%k_half1, %c48] shape [%k, %n] : view<64x64xf16, %activation_fragment_layout> -> vector<16xf16> + %half_next12 = vector.mma %lhs1, %rhs12, %half_acc12 : vector<16xf16>, vector<16xf16>, vector<8xf32> + %half_next13 = vector.mma %lhs1, %rhs13, %half_acc13 : vector<16xf16>, vector<16xf16>, vector<8xf32> + scf.schedule.fence + scf.yield %half_next10, %half_next11, %half_next12, %half_next13 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %next00, %next01, %next02, %next03, %next10, %next11, %next12, %next13 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + scf.yield %block_result00, %block_result01, %block_result02, %block_result03, %block_result10, %block_result11, %block_result12, %block_result13 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result00, %c0, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result01, %c16, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result02, %c32, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result03, %c48, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result10, %c64, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result11, %c80, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result12, %c96, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result13, %c112, %token_tile_base, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index) + template.return +} + +// PP128+ wave32 specialization for ordinary F32 activations. Keeping this as +// a separate provider leaves the wave64 decode schedule and ABI unchanged. +kernel.def target(@ggml_mul_mat_q6_k_prefill_gfx11_wave32) export("ggml_mul_mat_q6_k_f32_wmma_prefill_wave32") @ggml_mul_mat_q6_k_f32_wmma_prefill_wave32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %token_capacity = config.get @ggml.mul_mat_q6_k_packed.token_capacity : index + %c1 = index.constant 1 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %padded_output_size = index.add %output_size, %c127 : index + %output_tiles = index.div %padded_output_size, %c128 : index + %padded_token_count = index.add %token_capacity, %c127 : index + %token_tiles = index.div %padded_token_count, %c128 : index + kernel.launch.config workgroups(%token_tiles, %output_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_q6_k_packed.input_size : index + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %token_capacity = config.get @ggml.mul_mat_q6_k_packed.token_capacity : index + %bounded_token_count = index.assume %token_count [range(%token_count, 128, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %input_is_f16 = scalar.constant false : i1 + template.apply<@ggml.mul_mat_q6_k_packed.prefill_wave32.body>(%input_is_f16, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %weight, %output) : (i1, index, index, index, index, index, buffer, buffer, buffer) + kernel.return +} + +// PP128+ wave32 specialization for the F16 transient produced by the fused +// gate/up/SwiGLU recipe. +// Complete M512 tiles amortize the private F16 copy on sufficiently large grids. +func.def pure inline @ggml_q6k_prefill_packed_tile(%tokens: index, %inputs: index, %outputs: index, %accumulation: index) -> (i1, index) { + %zero = index.constant 0 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c512 = index.constant 512 : index + %contracting = index.cmp ugt, %inputs, %outputs : index + %n_tile = scf.select %contracting, %c64, %c128 : index + %m_tail = index.rem %tokens, %c512 : index + %n_tail = index.rem %outputs, %n_tile : index + %full_m = index.cmp eq, %m_tail, %zero : index + %full_n = index.cmp eq, %n_tail, %zero : index + %overwrite = index.cmp eq, %accumulation, %zero : index + %full_tile = scalar.andi %full_m, %full_n : i1 + %full_output = scalar.andi %full_tile, %overwrite : i1 + %m_tiles = index.div %tokens, %c512 : index + %n_tiles = index.div %outputs, %n_tile : index + %groups = index.mul %m_tiles, %n_tiles : index + %enough_groups = index.cmp uge, %groups, %c32 : index + %packed = scalar.andi %full_output, %enough_groups : i1 + func.return %packed, %n_tile : i1, index +} + +kernel.def target(@ggml_mul_mat_q6_k_prefill_gfx11_wave32) export("ggml_mul_mat_q6_k_f16_wmma_prefill_wave32") @ggml_mul_mat_q6_k_f16_wmma_prefill_wave32(%token_count: index) { + %input_size = config.get @ggml.mul_mat_q6_k_packed.input_size : index + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %token_capacity = config.get @ggml.mul_mat_q6_k_packed.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %c512 = index.constant 512 : index + %c1024 = index.constant 1024 : index + %minimum_workgroups = index.constant 32 : index + %token_tail = index.rem %token_count, %c256 : index + %channel_tail = index.rem %output_size, %c128 : index + %full_tokens = index.cmp eq, %token_tail, %c0 : index + %full_channels = index.cmp eq, %channel_tail, %c0 : index + %overwrite = index.cmp eq, %output_accumulation, %c0 : index + %full_tile = scalar.andi %full_tokens, %full_channels : i1 + %wide_tile = scalar.andi %full_tile, %overwrite : i1 + %padded_output_size = index.add %output_size, %c127 : index + %ordinary_output_tiles = index.div %padded_output_size, %c128 : index + %padded_token_count = index.add %token_capacity, %c127 : index + %small_token_tiles = index.div %padded_token_count, %c128 : index + %wide_token_tiles = index.div %token_count, %c256 : index + %staged_token_tiles = scf.select %wide_tile, %wide_token_tiles, %small_token_tiles : index + %wide_workgroup_size = scf.select %wide_tile, %c512, %c256 : index + %larger_output_tiles = index.div %output_size, %c256 : index + %larger_grid = index.mul %wide_token_tiles, %larger_output_tiles : index + %enough_workgroups = index.cmp uge, %larger_grid, %minimum_workgroups : index + %larger_channel_tail = index.rem %output_size, %c256 : index + %larger_full_channels = index.cmp eq, %larger_channel_tail, %c0 : index + %larger_full_tile = scalar.andi %wide_tile, %larger_full_channels : i1 + %multiple_token_tiles = index.cmp ugt, %wide_token_tiles, %c1 : index + %larger_grid_tile = scalar.andi %larger_full_tile, %enough_workgroups : i1 + %larger_tile = scalar.andi %larger_grid_tile, %multiple_token_tiles : i1 + %staged_workgroup_size = scf.select %larger_tile, %c1024, %wide_workgroup_size : index + %staged_output_tiles = scf.select %larger_tile, %larger_output_tiles, %ordinary_output_tiles : index + %packed, %packed_n_tile = func.call pure @ggml_q6k_prefill_packed_tile(%token_count, %input_size, %output_size, %output_accumulation) : (index, index, index, index) -> (i1, index) + %c8 = index.constant 8 : index + %packed_workgroup_size = index.mul %packed_n_tile, %c8 : index + %packed_token_tiles = index.div %token_count, %c512 : index + %packed_output_tiles = index.div %output_size, %packed_n_tile : index + %workgroup_size = scf.select %packed, %packed_workgroup_size, %staged_workgroup_size : index + %token_tiles = scf.select %packed, %packed_token_tiles, %staged_token_tiles : index + %output_tiles = scf.select %packed, %packed_output_tiles, %staged_output_tiles : index + kernel.launch.config workgroups(%token_tiles, %output_tiles, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) { + %input_size = config.get @ggml.mul_mat_q6_k_packed.input_size : index + %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index + %token_capacity = config.get @ggml.mul_mat_q6_k_packed.token_capacity : index + %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index + %packed_input, %packed_n_tile = func.call pure @ggml_q6k_prefill_packed_tile(%token_capacity, %input_size, %output_size, %output_accumulation) : (index, index, index, index) -> (i1, index) + %bounded_token_count = index.assume %token_count [range(%token_count, 128, 2048), le(%token_count, %token_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %workgroup_size = kernel.workgroup.size : index + %c512 = index.constant 512 : index + %wide_tile = index.cmp uge, %workgroup_size, %c512 : index + scf.if %wide_tile { + %weight_format = index.constant 6 : index + %paired = scalar.constant false : i1 + %prefill_binary_op = index.constant 0 : index + func.call @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format, %prefill_binary_op, %paired, %packed_input, %paired, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %weight, %weight, %output, %output) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + } else { + %input_is_f16 = scalar.constant true : i1 + template.apply<@ggml.mul_mat_q6_k_packed.prefill_wave32.body>(%input_is_f16, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %weight, %output) : (i1, index, index, index, index, index, buffer, buffer, buffer) + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_decode.loom new file mode 100644 index 000000000000..d0cc802ddcf8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_decode.loom @@ -0,0 +1,68 @@ +// Dense GGML decode gate/up matmul fused with an ordered binary op. +// +// The shared dual decode core computes the two adjacent gate/up row pairs. +// This file supplies the SwiGLU-specific publish hook and launcher. +template.decl @ggml.mul_mat_f32_f32_decode.dual_dispatch(%op: index, %lhs_weight_format: index, %rhs_weight_format: index, %token_capacity: index, %token_count: index, %input_size: index, %output_size: index, %input: buffer, %lhs_weight: buffer, %rhs_weight: buffer, %output: buffer) + +template.decl @ggml.mul_mat_f32_f32_decode.publish_dual_pair2(%op: index, %publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %lhs0: f32, %lhs1: f32, %rhs0: f32, %rhs1: f32, %output: buffer) + +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_mul_mat_swiglu_decode_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 1)] + +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)] + +template.def<@ggml.mul_mat_f32_f32_decode.publish_dual_pair2> device priority(20) @ggml_mul_mat_swiglu_decode_publish_dual_pair2(%op: index, %publish_pair: i1, %token0: index, %channel: index, %token_count0: index, %output_size0: index, %row1_valid: i1, %gate0: f32, %gate1: f32, %up0: f32, %up1: f32, %output: buffer) { + %token_count = index.assume %token_count0 [range(%token_count0, 1, 2048)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 1, 1048576)] : index + %c1 = index.constant 1 : index + %c0_offset = index.constant 0 : offset + %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index + %output_noalias = buffer.assume.noalias %output : buffer + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%output_size]xf32> + scf.if %publish_pair { + %value0 = template.apply<@ggml.binary_f32.apply>(%op, %gate0, %up0) : (index, f32, f32) -> (f32) + view.store %value0, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + scf.if %row1_valid { + scf.if %publish_pair { + %row1 = index.add %channel, %c1 : index + %value1 = template.apply<@ggml.binary_f32.apply>(%op, %gate1, %up1) : (index, f32, f32) -> (f32) + view.store %value1, %output_view[%token, %row1] : f32, view<[%launch_token_count]x[%output_size]xf32> + } + } + template.return +} + +kernel.def target(@ggml_mul_mat_swiglu_decode_gfx11_wave64) @ggml_mul_mat_swiglu_f32_f32_decode_wave64(%token_count: index) { + %token_capacity = config.get @ggml.workload.token_capacity : index + %output_capacity = config.get @ggml.mul_mat_swiglu.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %padded_output_size = index.add %output_capacity, %c1 : index + %output_pairs = index.div %padded_output_size, %c2 : index + %padded_token_count = index.add %token_capacity, %c63 : index + %launch_tokens = index.div %padded_token_count, %c64 : index + kernel.launch.config workgroups(%output_pairs, %launch_tokens, %c1) workgroup_size(%c64, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 1)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : 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 + %launch_token_capacity = config.get @ggml.workload.token_capacity : index + template.apply<@ggml.mul_mat_f32_f32_decode.dual_dispatch>(%op, %gate_weight_format, %up_weight_format, %launch_token_capacity, %token_count, %input_size, %output_size, %input, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom new file mode 100644 index 000000000000..f81e0f3f6593 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_f32_f32_wmma.loom @@ -0,0 +1,1558 @@ +// Dense GGML gate/up matmul fused with an ordered binary op for contiguous 2D F32 activations. +func.decl @ggml_iq4xs_load_decode_state(%weight: buffer, %row_base: offset, %block: index) -> (i32, i32) + +func.decl @ggml_iq4xs_preloaded_f32_vector4(%iq4nl_table: vector<16xi8>, %weight: buffer, %row_byte_base: offset, %iq4_block: index, %iq4_group: index, %packet: index, %header: i32, %scales_l: i32) -> (vector<4xf32>) + +func.decl @ggml_q5k_load_decode_state(%weight: buffer, %row_base: offset, %block: index, %packet: index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + +func.decl @ggml_q5k_preloaded_f32_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>, %qh_bytes: vector<4xi8>) -> (vector<4xf32>) + +template.decl @ggml.mul_mat_swiglu.apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) + +func.decl @ggml_q4k_scale_min_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) + +func.decl @ggml_dot_u8_s8_vector8_f32(%weight: vector<2xi32>, %activation: vector<2xi32>) -> (f32) + +func.decl @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf32>) + +func.decl @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf32>) + +template.decl @ggml.unary_f32.apply_vector4(%op: index, %values: vector<4xf32>) -> (vector<4xf32>) + +template.decl @ggml.mul_mat_swiglu.body(%op: 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, %output: buffer) + +template.decl @ggml.mul_mat_swiglu.launch(%token_capacity: index, %output_size: index) -> (index, index, index, index) + +amdgpu.target @ggml_mul_mat_swiglu_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.mul_mat_swiglu.f16_output_layout : %value: index where [range(%value, 0, 1)] + +config.decl @ggml.workload.token_capacity : %value: index where [range(%value, 1, 2048)] + +func.decl @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) + +func.decl @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offset) + +func.decl @ggml_iq4nl_table_i8() -> (vector<16xi8>) + +func.decl @ggml_iq3s_grid_f32_halves() -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + +func.decl @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_grid_low0: vector<32xf32>, %iq3s_grid_low1: vector<32xf32>, %iq3s_grid_low2: vector<32xf32>, %iq3s_grid_low3: vector<32xf32>, %iq3s_grid_low4: vector<32xf32>, %iq3s_grid_low5: vector<32xf32>, %iq3s_grid_low6: vector<32xf32>, %iq3s_grid_low7: vector<32xf32>, %iq3s_grid_low8: vector<32xf32>, %iq3s_grid_low9: vector<32xf32>, %iq3s_grid_low10: vector<32xf32>, %iq3s_grid_low11: vector<32xf32>, %iq3s_grid_low12: vector<32xf32>, %iq3s_grid_low13: vector<32xf32>, %iq3s_grid_low14: vector<32xf32>, %iq3s_grid_low15: vector<32xf32>, %iq3s_grid_high0: vector<32xf32>, %iq3s_grid_high1: vector<32xf32>, %iq3s_grid_high2: vector<32xf32>, %iq3s_grid_high3: vector<32xf32>, %iq3s_grid_high4: vector<32xf32>, %iq3s_grid_high5: vector<32xf32>, %iq3s_grid_high6: vector<32xf32>, %iq3s_grid_high7: vector<32xf32>, %iq3s_grid_high8: vector<32xf32>, %iq3s_grid_high9: vector<32xf32>, %iq3s_grid_high10: vector<32xf32>, %iq3s_grid_high11: vector<32xf32>, %iq3s_grid_high12: vector<32xf32>, %iq3s_grid_high13: vector<32xf32>, %iq3s_grid_high14: vector<32xf32>, %iq3s_grid_high15: vector<32xf32>, %weight_format: index, %weight: buffer, %row_byte_base: offset, %input_size: index, %quant_block: index, %quant_group: index, %packet: index, %k: index) -> (vector<4xf16>) + +template.def<@ggml.mul_mat_swiglu.apply_vector4> device @ggml_mul_mat_swiglu_apply_vector4(%op: index, %lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (vector<4xf32>) { + %c0_5_f32x4 = vector.constant 0.5 : vector<4xf32> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c1_f32x4 = vector.constant 1.0 : vector<4xf32> + %c2_f32x4 = vector.constant 2.0 : vector<4xf32> + %cn2_f32x4 = vector.constant -2.0 : vector<4xf32> + %gelu_coef = vector.constant 0.044715 : vector<4xf32> + %sqrt_2_over_pi = vector.constant 0.79788456080286544 : vector<4xf32> + + %op_add = index.constant 0 : index + %sum = vector.addf %lhs, %rhs : vector<4xf32> + %is_add = index.cmp eq, %op, %op_add : index + %add_selected = scf.select %is_add, %sum, %lhs : vector<4xf32> + + %op_sub = index.constant 1 : index + %difference = vector.subf %lhs, %rhs : vector<4xf32> + %is_sub = index.cmp eq, %op, %op_sub : index + %sub_selected = scf.select %is_sub, %difference, %add_selected : vector<4xf32> + + %op_mul = index.constant 2 : index + %product = vector.mulf %lhs, %rhs : vector<4xf32> + %is_mul = index.cmp eq, %op, %op_mul : index + %mul_selected = scf.select %is_mul, %product, %sub_selected : vector<4xf32> + + %op_div = index.constant 3 : index + %quotient = vector.divf %lhs, %rhs : vector<4xf32> + %is_div = index.cmp eq, %op, %op_div : index + %div_selected = scf.select %is_div, %quotient, %mul_selected : vector<4xf32> + + %op_swiglu = index.constant 4 : index + %negative_lhs = vector.subf %c0_f32x4, %lhs : vector<4xf32> + %lhs_exp = vector.expf %negative_lhs : vector<4xf32> + %denominator = vector.addf %c1_f32x4, %lhs_exp : vector<4xf32> + %sigmoid = vector.divf %c1_f32x4, %denominator : vector<4xf32> + %silu = vector.mulf %lhs, %sigmoid : vector<4xf32> + %swiglu = vector.mulf %silu, %rhs : vector<4xf32> + %is_swiglu = index.cmp eq, %op, %op_swiglu : index + %swiglu_selected = scf.select %is_swiglu, %swiglu, %div_selected : vector<4xf32> + + %op_geglu = index.constant 5 : index + %gelu_x2 = vector.mulf %lhs, %lhs : vector<4xf32> + %gelu_poly0 = vector.mulf %gelu_coef, %gelu_x2 : vector<4xf32> + %gelu_poly = vector.addf %c1_f32x4, %gelu_poly0 : vector<4xf32> + %gelu_inner0 = vector.mulf %lhs, %gelu_poly : vector<4xf32> + %gelu_inner = vector.mulf %sqrt_2_over_pi, %gelu_inner0 : vector<4xf32> + %gelu_tanh_scaled = vector.mulf %cn2_f32x4, %gelu_inner : vector<4xf32> + %gelu_tanh_exp = vector.expf %gelu_tanh_scaled : vector<4xf32> + %gelu_tanh_denominator = vector.addf %c1_f32x4, %gelu_tanh_exp : vector<4xf32> + %gelu_tanh_ratio = vector.divf %c2_f32x4, %gelu_tanh_denominator : vector<4xf32> + %gelu_tanh = vector.subf %gelu_tanh_ratio, %c1_f32x4 : vector<4xf32> + %gelu_one_plus = vector.addf %c1_f32x4, %gelu_tanh : vector<4xf32> + %gelu_half_x = vector.mulf %c0_5_f32x4, %lhs : vector<4xf32> + %gelu = vector.mulf %gelu_half_x, %gelu_one_plus : vector<4xf32> + %geglu = vector.mulf %gelu, %rhs : vector<4xf32> + %is_geglu = index.cmp eq, %op, %op_geglu : index + %geglu_selected = scf.select %is_geglu, %geglu, %swiglu_selected : vector<4xf32> + + %op_reglu = index.constant 6 : index + %unary_relu = index.constant 2 : index + %relu_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_relu, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %reglu = vector.mulf %relu_lhs, %rhs : vector<4xf32> + %is_reglu = index.cmp eq, %op, %op_reglu : index + %reglu_selected = scf.select %is_reglu, %reglu, %geglu_selected : vector<4xf32> + + %op_geglu_erf = index.constant 7 : index + %unary_gelu_erf = index.constant 20 : index + %gelu_erf_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_gelu_erf, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %geglu_erf = vector.mulf %gelu_erf_lhs, %rhs : vector<4xf32> + %is_geglu_erf = index.cmp eq, %op, %op_geglu_erf : index + %geglu_erf_selected = scf.select %is_geglu_erf, %geglu_erf, %reglu_selected : vector<4xf32> + + %op_geglu_quick = index.constant 8 : index + %unary_gelu_quick = index.constant 14 : index + %gelu_quick_lhs = template.apply<@ggml.unary_f32.apply_vector4>(%unary_gelu_quick, %lhs) : (index, vector<4xf32>) -> (vector<4xf32>) + %geglu_quick = vector.mulf %gelu_quick_lhs, %rhs : vector<4xf32> + %is_geglu_quick = index.cmp eq, %op, %op_geglu_quick : index + %result = scf.select %is_geglu_quick, %geglu_quick, %geglu_erf_selected : vector<4xf32> + template.return %result : vector<4xf32> +} + +template.def<@ggml.mul_mat_swiglu.body> device @ggml_mul_mat_swiglu_f32_f32_wmma_body(%op: 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, %output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 2, 2048)] : index + %bounded_input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 32)] : index + %bounded_output_size = index.assume %output_size0 [range(%output_size0, 1, 262144)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c40 = index.constant 40 : index + %c64 = index.constant 64 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %c0_offset = index.constant 0 : offset + %weight_stage_bytes = index.constant 5120 : offset + %activation_stage_bytes = index.constant 2560 : offset + %wave_result_stage_bytes = index.constant 512 : offset + %result_stage_bytes = index.constant 1024 : offset + %c0_f16x4 = vector.constant 0.0 : vector<4xf16> + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %zero_accumulator = vector.constant 0.0 : vector<8xf16> + %m = index.constant 16 : index + %n = index.constant 16 : index + %k = index.constant 16 : index + %padded_input_size = index.add %bounded_input_size, %c255 : index + %quant_block_count = index.div %padded_input_size, %c256 : index + %gate_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%gate_weight_format, %bounded_input_size) : (index, index) -> (offset) + %up_weight_row_bytes = func.call @ggml_dequant_weight_row_bytes(%up_weight_format, %bounded_input_size) : (index, index) -> (offset) + %input_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> + %gate_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %up_weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer + %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer + %gate_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %up_result_stage = buffer.alloca align(16) %result_stage_bytes : buffer + %gate_weight_stage_view = buffer.view %gate_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %up_weight_stage_view = buffer.view %up_weight_stage[%c0_offset] : buffer -> view<64x40xf16> + %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> + %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding + %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> + %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset + %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding + %gate_result_fragment_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %up_result_fragment_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> + %gate_result_physical_view = buffer.view %gate_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %up_result_physical_view = buffer.view %up_result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> + %channel_tile_base = index.mul %channel_tile, %c64 : index + %token_tile_base = index.mul %token_tile, %c32 : index + %load_packet = index.rem %workitem, %c8 : index + %load_k = index.mul %load_packet, %c4 : index + %load_row0 = index.div %workitem, %c8 : index + %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index + %subgroup_channel_add = index.mul %subgroup, %c32 : index + %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index + %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %gate_result00, %gate_result01, %gate_result10, %gate_result11, %up_result00, %up_result01, %up_result10, %up_result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%gate_block_acc00 = %init00 : vector<8xf16>, %gate_block_acc01 = %init01 : vector<8xf16>, %gate_block_acc10 = %init10 : vector<8xf16>, %gate_block_acc11 = %init11 : vector<8xf16>, %up_block_acc00 = %init00 : vector<8xf16>, %up_block_acc01 = %init01 : vector<8xf16>, %up_block_acc10 = %init10 : vector<8xf16>, %up_block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%gate_acc00 = %gate_block_acc00 : vector<8xf16>, %gate_acc01 = %gate_block_acc01 : vector<8xf16>, %gate_acc10 = %gate_block_acc10 : vector<8xf16>, %gate_acc11 = %gate_block_acc11 : vector<8xf16>, %up_acc00 = %up_block_acc00 : vector<8xf16>, %up_acc01 = %up_block_acc01 : vector<8xf16>, %up_acc10 = %up_block_acc10 : vector<8xf16>, %up_acc11 = %up_block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { + %block_k_base = index.mul %quant_block, %c256 : index + %group_k_add = index.mul %quant_group, %c32 : index + %k_origin = index.add %block_k_base, %group_k_add : index + scf.for %row_offset = [%c0 to %c64 step %c16] unroll { + %local_row0 = index.add %load_row, %row_offset : index + %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index + %channel = index.add %channel_tile_base, %local_row : index + %valid_channel = index.cmp ult, %channel, %bounded_output_size : index + %weight_k = index.add %k_origin, %load_k : index + %valid_k = index.cmp ult, %weight_k, %bounded_input_size : index + %valid_weight = scalar.andi %valid_channel, %valid_k : i1 + %gate_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %gate_weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + %up_weight_values = scf.if %valid_weight -> (vector<4xf16>) { + %row_byte_base = index.scale %channel, %up_weight_row_bytes : index, offset -> offset + %decoded = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %decoded : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %gate_weight_values, %gate_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + vector.store %up_weight_values, %up_weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> + %is_activation_row = index.cmp ult, %local_row, %c32 : index + scf.if %is_activation_row { + %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index + %token = index.add %token_tile_base, %activation_row : index + %valid_token = index.cmp ult, %token, %bounded_token_count : index + %input_k = index.add %k_origin, %load_k : index + %valid_input = index.cmp ult, %input_k, %bounded_input_size : index + %valid_activation = scalar.andi %valid_token, %valid_input : i1 + %activation_values = scf.if %valid_activation -> (vector<4xf16>) { + %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index + %mask = vector.mask.range [%input_k to %bounded_input_size step %c1] : index -> vector<4xi1> + %loaded = vector.load.mask %input_view[%bounded_token, %input_k], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_input_size]xf32>, vector<4xi1>, vector<4xf32> + %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> + scf.yield %converted : vector<4xf16> + } else { + scf.yield %c0_f16x4 : vector<4xf16> + } + vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 = scf.for %k_half = [%c0 to %c32 step %c16](%gate_half_acc00 = %gate_acc00 : vector<8xf16>, %gate_half_acc01 = %gate_acc01 : vector<8xf16>, %gate_half_acc10 = %gate_acc10 : vector<8xf16>, %gate_half_acc11 = %gate_acc11 : vector<8xf16>, %up_half_acc00 = %up_acc00 : vector<8xf16>, %up_half_acc01 = %up_acc01 : vector<8xf16>, %up_half_acc10 = %up_acc10 : vector<8xf16>, %up_half_acc11 = %up_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { + %gate_lhs0 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %gate_lhs1 = vector.fragment.load %gate_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs0 = vector.fragment.load %up_weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %up_lhs1 = vector.fragment.load %up_weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> + %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> + %gate_half_next00 = vector.mma %gate_lhs0, %rhs0, %gate_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next01 = vector.mma %gate_lhs0, %rhs1, %gate_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next10 = vector.mma %gate_lhs1, %rhs0, %gate_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %gate_half_next11 = vector.mma %gate_lhs1, %rhs1, %gate_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next00 = vector.mma %up_lhs0, %rhs0, %up_half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next01 = vector.mma %up_lhs0, %rhs1, %up_half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next10 = vector.mma %up_lhs1, %rhs0, %up_half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> + %up_half_next11 = vector.mma %up_lhs1, %rhs1, %up_half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> + scf.yield %gate_half_next00, %gate_half_next01, %gate_half_next10, %gate_half_next11, %up_half_next00, %up_half_next01, %up_half_next10, %up_half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %gate_next00, %gate_next01, %gate_next10, %gate_next11, %up_next00, %up_next01, %up_next10, %up_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + scf.yield %gate_block_result00, %gate_block_result01, %gate_block_result10, %gate_block_result11, %up_block_result00, %up_block_result01, %up_block_result10, %up_block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> + } + %publish_token0 = index.div %lane, %c4 : index + %publish_token = index.assume %publish_token0 [range(%publish_token0, 0, 15)] : index + %publish_packet0 = index.rem %lane, %c4 : index + %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index + %publish_channel_add = index.mul %publish_packet, %c4 : index + %token0 = index.add %token_tile_base, %publish_token : index + %publish_token1 = index.add %publish_token, %c16 : index + %token1 = index.add %token_tile_base, %publish_token1 : index + %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index + %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index + %channel1_base = index.add %subgroup_channel_base, %c16 : index + %channel1 = index.add %channel1_base, %publish_channel_add : index + %valid_token0 = index.cmp ult, %token0, %bounded_token_count : index + %valid_token1 = index.cmp ult, %token1, %bounded_token_count : index + %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index + %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index + %writes00 = scalar.andi %valid_token0, %valid_channel0 : i1 + %writes01 = scalar.andi %valid_token1, %valid_channel0 : i1 + %writes10 = scalar.andi %valid_token0, %valid_channel1 : i1 + %writes11 = scalar.andi %valid_token1, %valid_channel1 : i1 + vector.fragment.store %gate_result00, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result00, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes00 { + %bounded_token, %output_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gate_wide, %up_wide) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_token, %channel0], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result01, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result01, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes01 { + %bounded_token, %output_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gate_wide, %up_wide) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_token, %channel0], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result10, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result10, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes10 { + %bounded_token, %output_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gate_wide, %up_wide) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_token, %channel1], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + kernel.barrier scope(subgroup) ordering(acq_rel) + vector.fragment.store %gate_result11, %gate_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + vector.fragment.store %up_result11, %up_result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> + kernel.barrier scope(subgroup) ordering(acq_rel) + scf.if %writes11 { + %bounded_token, %output_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index + %gate_values = vector.load %gate_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %up_values = vector.load %up_result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> + %gate_wide = vector.extf %gate_values : vector<4xf16> to vector<4xf32> + %up_wide = vector.extf %up_values : vector<4xf16> to vector<4xf32> + %activated = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gate_wide, %up_wide) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> + vector.store.mask %activated, %output_view[%bounded_token, %channel1], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> + } + template.return +} + +template.def<@ggml.mul_mat_swiglu.launch> @ggml_mul_mat_swiglu_f32_f32_wmma_launch(%token_capacity: index, %output_size: index) -> (index, index, index, 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 + template.return %output_tiles, %token_tiles, %c1, %c128 : index, index, index, index +} + +kernel.def target(@ggml_mul_mat_swiglu_gfx11_wave64) @ggml_mul_mat_swiglu_f32_f32_wmma(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %c63 = index.constant 63 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %padded_output_size = index.add %output_size, %c63 : index + %output_tiles = index.div %padded_output_size, %c64 : index + %padded_token_count = index.add %token_capacity, %c31 : index + %token_tiles = index.div %padded_token_count, %c32 : index + kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 2, 2048)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %gate_weight_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_weight_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %op = config.get @ggml.mul_mat_swiglu.op : index + %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.mul_mat_swiglu.body>(%op, %gate_weight_format, %up_weight_format, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +amdgpu.target @ggml_mul_mat_swiglu_gfx11_wave32 {subgroup_size = 32} + +// Low-token specializations keep FP32 operands and accumulation and avoid +// padding a handful of tokens to a full matrix tile. +// Share the activation load across gate/up while preserving token guards and FP32 update order. +func.def inline @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens: index, %k: index, %input: view<[%tokens]x[%k]xf32>, %token: index, %kk: index, %gate: vector<4xf32>, %up: vector<4xf32>, %acc_gate: vector<4xf32>, %acc_up: vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) { + %valid = index.cmp ult, %token, %tokens : index + %next_gate, %next_up = scf.if %valid -> (vector<4xf32>, vector<4xf32>) { + %a = vector.load %input[%token, %kk] : view<[%tokens]x[%k]xf32> -> vector<4xf32> + %gr = vector.fmaf %gate, %a, %acc_gate : vector<4xf32> + %ur = vector.fmaf %up, %a, %acc_up : vector<4xf32> + scf.yield %gr, %ur : vector<4xf32>, vector<4xf32> + } else { + scf.yield %acc_gate, %acc_up : vector<4xf32>, vector<4xf32> + } + func.return %next_gate, %next_up : vector<4xf32>, vector<4xf32> +} + +func.def inline @ggml_mul_mat_swiglu_lowtoken_dot(%op: index, %gate_format: index, %up_format: index, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %input: buffer, %gate: buffer, %up: buffer, %output: buffer) { + %iq4nl_table = func.call @ggml_iq4nl_table_i8() : () -> (vector<16xi8>) + %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15 = func.call @ggml_iq3s_grid_f32_halves() : () -> (vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>) + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c10 = index.constant 10 : index + %c256 = index.constant 256 : index + %q5_format = index.constant 5 : index + %iq4xs_format = index.constant 23 : index + %q8_0_format = index.constant 80 : index + %gate_is_q5 = index.cmp eq, %gate_format, %q5_format : index + %gate_is_iq4xs = index.cmp eq, %gate_format, %iq4xs_format : index + %same_weight_format = index.cmp eq, %gate_format, %up_format : index + %gate_is_q8_0 = index.cmp eq, %gate_format, %q8_0_format : index + %q8_0_pair = scalar.andi %gate_is_q8_0, %same_weight_format : i1 + %token_one = index.cmp eq, %token_count, %c1 : index + %q8_0_token_one = scalar.andi %q8_0_pair, %token_one : i1 + %fallback_unroll = scf.select %q8_0_token_one, %c10, %c1 : index + %base = index.constant 0 : offset + %zero = vector.constant 0.0 : vector<4xf32> + %zero_scalar = scalar.constant 0.0 : f32 + %tokens = index.assume %token_count [range(%token_count, 1, 5)] : index + %k = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index + %n = index.assume %output_size [range(%output_size, 1, 262144)] : index + %subgroup = kernel.subgroup.id : index + %lane = kernel.subgroup.lane.id : index + %c16 = index.constant 16 : index + %c64 = index.constant 64 : index + %row_base = index.mul %channel_tile, %c8 : index + %wave_row = index.mul %subgroup, %c2 : index + %cohort = index.div %lane, %c16 : index + %local_row = index.add %wave_row, %cohort : index + %row = index.add %row_base, %local_row : index + %lane16 = index.rem %lane, %c16 : index + %valid_row = index.cmp ult, %row, %n : index + %lane_group = index.div %lane16, %c8 : index + %packet = index.rem %lane16, %c8 : index + %lane_k = index.mul %lane16, %c4 : index + %blocks = index.div %k, %c256 : index + %steps = index.div %k, %c64 : index + %gate_bytes = func.call @ggml_dequant_weight_tile_bytes(%gate_format) : (index) -> (offset) + %up_bytes = func.call @ggml_dequant_weight_tile_bytes(%up_format) : (index) -> (offset) + %gate_row_bytes = index.scale %blocks, %gate_bytes : index, offset -> offset + %up_row_bytes = index.scale %blocks, %up_bytes : index, offset -> offset + %gate_base = index.scale %row, %gate_row_bytes : index, offset -> offset + %up_base = index.scale %row, %up_row_bytes : index, offset -> offset + %a_na, %g_na, %u_na, %o_na = buffer.assume.noalias %input, %gate, %up, %output : buffer, buffer, buffer, buffer + %av = buffer.view %a_na[%base] : buffer -> view<[%tokens]x[%k]xf32> + %ov = buffer.view %o_na[%base] : buffer -> view<[%tokens]x[%n]xf32> + scf.if %valid_row { + // Expose one quantization block at a time so its four quarters can share scales and address arithmetic. + %q5_pair = scalar.andi %gate_is_q5, %same_weight_format : i1 + %total_g0, %total_u0, %total_g1, %total_u1, %total_g2, %total_u2, %total_g3, %total_u3, %total_g4, %total_u4 = scf.if %q5_pair -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %branch_total_g0, %branch_total_u0, %branch_total_g1, %branch_total_u1, %branch_total_g2, %branch_total_u2, %branch_total_g3, %branch_total_u3, %branch_total_g4, %branch_total_u4 = scf.for %block = [%c0 to %blocks step %c1](%block_acc_g0 = %zero : vector<4xf32>, %block_acc_u0 = %zero : vector<4xf32>, %block_acc_g1 = %zero : vector<4xf32>, %block_acc_u1 = %zero : vector<4xf32>, %block_acc_g2 = %zero : vector<4xf32>, %block_acc_u2 = %zero : vector<4xf32>, %block_acc_g3 = %zero : vector<4xf32>, %block_acc_u3 = %zero : vector<4xf32>, %block_acc_g4 = %zero : vector<4xf32>, %block_acc_u4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %g_dm, %g_scales, %g_qh = func.call @ggml_q5k_load_decode_state(%g_na, %gate_base, %block, %packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + %u_dm, %u_scales, %u_qh = func.call @ggml_q5k_load_decode_state(%u_na, %up_base, %block, %packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + %quarter_total_g0, %quarter_total_u0, %quarter_total_g1, %quarter_total_u1, %quarter_total_g2, %quarter_total_u2, %quarter_total_g3, %quarter_total_u3, %quarter_total_g4, %quarter_total_u4 = scf.for %quarter = [%c0 to %c4 step %c1](%acc_g0 = %block_acc_g0 : vector<4xf32>, %acc_u0 = %block_acc_u0 : vector<4xf32>, %acc_g1 = %block_acc_g1 : vector<4xf32>, %acc_u1 = %block_acc_u1 : vector<4xf32>, %acc_g2 = %block_acc_g2 : vector<4xf32>, %acc_u2 = %block_acc_u2 : vector<4xf32>, %acc_g3 = %block_acc_g3 : vector<4xf32>, %acc_u3 = %block_acc_u3 : vector<4xf32>, %acc_g4 = %block_acc_g4 : vector<4xf32>, %acc_u4 = %block_acc_u4 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) unroll(%c4) { + %step_base = index.mul %block, %c4 : index + %step = index.add %step_base, %quarter : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %block_k = index.mul %step, %c64 : index + %kk = index.add %block_k, %lane_k : index + %gh = func.call @ggml_q5k_preloaded_f32_vector4(%g_na, %gate_base, %block, %group, %packet, %g_dm, %g_scales, %g_qh) : (buffer, offset, index, index, index, vector<2xf16>, vector<3xi32>, vector<4xi8>) -> (vector<4xf32>) + %uh = func.call @ggml_q5k_preloaded_f32_vector4(%u_na, %up_base, %block, %group, %packet, %u_dm, %u_scales, %u_qh) : (buffer, offset, index, index, index, vector<2xf16>, vector<3xi32>, vector<4xi8>) -> (vector<4xf32>) + %t0 = index.constant 0 : index + %next_g0, %next_u0 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t0, %kk, %gh, %uh, %acc_g0, %acc_u0) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t1 = index.constant 1 : index + %next_g1, %next_u1 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t1, %kk, %gh, %uh, %acc_g1, %acc_u1) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t2 = index.constant 2 : index + %next_g2, %next_u2 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t2, %kk, %gh, %uh, %acc_g2, %acc_u2) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t3 = index.constant 3 : index + %next_g3, %next_u3 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t3, %kk, %gh, %uh, %acc_g3, %acc_u3) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t4 = index.constant 4 : index + %next_g4, %next_u4 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t4, %kk, %gh, %uh, %acc_g4, %acc_u4) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + scf.yield %next_g0, %next_u0, %next_g1, %next_u1, %next_g2, %next_u2, %next_g3, %next_u3, %next_g4, %next_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %quarter_total_g0, %quarter_total_u0, %quarter_total_g1, %quarter_total_u1, %quarter_total_g2, %quarter_total_u2, %quarter_total_g3, %quarter_total_u3, %quarter_total_g4, %quarter_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %branch_total_g0, %branch_total_u0, %branch_total_g1, %branch_total_u1, %branch_total_g2, %branch_total_u2, %branch_total_g3, %branch_total_u3, %branch_total_g4, %branch_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } else { + %iq4_pair = scalar.andi %gate_is_iq4xs, %same_weight_format : i1 + %branch_total_g0, %branch_total_u0, %branch_total_g1, %branch_total_u1, %branch_total_g2, %branch_total_u2, %branch_total_g3, %branch_total_u3, %branch_total_g4, %branch_total_u4 = scf.if %iq4_pair -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %iq4_total_g0, %iq4_total_u0, %iq4_total_g1, %iq4_total_u1, %iq4_total_g2, %iq4_total_u2, %iq4_total_g3, %iq4_total_u3, %iq4_total_g4, %iq4_total_u4 = scf.for %block = [%c0 to %blocks step %c1](%block_acc_g0 = %zero : vector<4xf32>, %block_acc_u0 = %zero : vector<4xf32>, %block_acc_g1 = %zero : vector<4xf32>, %block_acc_u1 = %zero : vector<4xf32>, %block_acc_g2 = %zero : vector<4xf32>, %block_acc_u2 = %zero : vector<4xf32>, %block_acc_g3 = %zero : vector<4xf32>, %block_acc_u3 = %zero : vector<4xf32>, %block_acc_g4 = %zero : vector<4xf32>, %block_acc_u4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { + %g_header, %g_scales_l = func.call @ggml_iq4xs_load_decode_state(%g_na, %gate_base, %block) : (buffer, offset, index) -> (i32, i32) + %u_header, %u_scales_l = func.call @ggml_iq4xs_load_decode_state(%u_na, %up_base, %block) : (buffer, offset, index) -> (i32, i32) + %quarter_total_g0, %quarter_total_u0, %quarter_total_g1, %quarter_total_u1, %quarter_total_g2, %quarter_total_u2, %quarter_total_g3, %quarter_total_u3, %quarter_total_g4, %quarter_total_u4 = scf.for %quarter = [%c0 to %c4 step %c1](%acc_g0 = %block_acc_g0 : vector<4xf32>, %acc_u0 = %block_acc_u0 : vector<4xf32>, %acc_g1 = %block_acc_g1 : vector<4xf32>, %acc_u1 = %block_acc_u1 : vector<4xf32>, %acc_g2 = %block_acc_g2 : vector<4xf32>, %acc_u2 = %block_acc_u2 : vector<4xf32>, %acc_g3 = %block_acc_g3 : vector<4xf32>, %acc_u3 = %block_acc_u3 : vector<4xf32>, %acc_g4 = %block_acc_g4 : vector<4xf32>, %acc_u4 = %block_acc_u4 : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) unroll(%c4) { + %step_base = index.mul %block, %c4 : index + %step = index.add %step_base, %quarter : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %block_k = index.mul %step, %c64 : index + %kk = index.add %block_k, %lane_k : index + %gh = func.call @ggml_iq4xs_preloaded_f32_vector4(%iq4nl_table, %g_na, %gate_base, %block, %group, %packet, %g_header, %g_scales_l) : (vector<16xi8>, buffer, offset, index, index, index, i32, i32) -> (vector<4xf32>) + %uh = func.call @ggml_iq4xs_preloaded_f32_vector4(%iq4nl_table, %u_na, %up_base, %block, %group, %packet, %u_header, %u_scales_l) : (vector<16xi8>, buffer, offset, index, index, index, i32, i32) -> (vector<4xf32>) + %t0 = index.constant 0 : index + %next_g0, %next_u0 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t0, %kk, %gh, %uh, %acc_g0, %acc_u0) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t1 = index.constant 1 : index + %next_g1, %next_u1 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t1, %kk, %gh, %uh, %acc_g1, %acc_u1) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t2 = index.constant 2 : index + %next_g2, %next_u2 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t2, %kk, %gh, %uh, %acc_g2, %acc_u2) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t3 = index.constant 3 : index + %next_g3, %next_u3 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t3, %kk, %gh, %uh, %acc_g3, %acc_u3) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t4 = index.constant 4 : index + %next_g4, %next_u4 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t4, %kk, %gh, %uh, %acc_g4, %acc_u4) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + scf.yield %next_g0, %next_u0, %next_g1, %next_u1, %next_g2, %next_u2, %next_g3, %next_u3, %next_g4, %next_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %quarter_total_g0, %quarter_total_u0, %quarter_total_g1, %quarter_total_u1, %quarter_total_g2, %quarter_total_u2, %quarter_total_g3, %quarter_total_u3, %quarter_total_g4, %quarter_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %iq4_total_g0, %iq4_total_u0, %iq4_total_g1, %iq4_total_u1, %iq4_total_g2, %iq4_total_u2, %iq4_total_g3, %iq4_total_u3, %iq4_total_g4, %iq4_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } else { + %fallback_total_g0, %fallback_total_u0, %fallback_total_g1, %fallback_total_u1, %fallback_total_g2, %fallback_total_u2, %fallback_total_g3, %fallback_total_u3, %fallback_total_g4, %fallback_total_u4 = scf.for %step = [%c0 to %steps step %c1](%acc_g0 = %zero : vector<4xf32>, %acc_u0 = %zero : vector<4xf32>, %acc_g1 = %zero : vector<4xf32>, %acc_u1 = %zero : vector<4xf32>, %acc_g2 = %zero : vector<4xf32>, %acc_u2 = %zero : vector<4xf32>, %acc_g3 = %zero : vector<4xf32>, %acc_u3 = %zero : vector<4xf32>, %acc_g4 = %zero : vector<4xf32>, %acc_u4 = %zero : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) unroll(%fallback_unroll) { + %block = index.div %step, %c4 : index + %quarter = index.rem %step, %c4 : index + %quarter_group = index.mul %quarter, %c2 : index + %group = index.add %quarter_group, %lane_group : index + %block_k = index.mul %step, %c64 : index + %kk = index.add %block_k, %lane_k : index + %g_packed_format = index.constant 44 : index + %g_packed = index.cmp eq, %gate_format, %g_packed_format : index + %gh = scf.if %g_packed -> (vector<4xf32>) { + %values = func.call @ggml_q4k_native_row64_f32_vector4(%g_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %values = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_format, %g_na, %gate_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } + %u_packed_format = index.constant 44 : index + %u_packed = index.cmp eq, %up_format, %u_packed_format : index + %uh = scf.if %u_packed -> (vector<4xf32>) { + %values = func.call @ggml_q4k_native_row64_f32_vector4(%u_na, %row, %k, %block, %group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %values = func.call @ggml_dequant_f32_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_format, %u_na, %up_base, %k, %block, %group, %packet, %kk) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } + %t0 = index.constant 0 : index + %next_g0, %next_u0 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t0, %kk, %gh, %uh, %acc_g0, %acc_u0) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t1 = index.constant 1 : index + %next_g1, %next_u1 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t1, %kk, %gh, %uh, %acc_g1, %acc_u1) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t2 = index.constant 2 : index + %next_g2, %next_u2 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t2, %kk, %gh, %uh, %acc_g2, %acc_u2) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t3 = index.constant 3 : index + %next_g3, %next_u3 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t3, %kk, %gh, %uh, %acc_g3, %acc_u3) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + %t4 = index.constant 4 : index + %next_g4, %next_u4 = func.call @ggml_mul_mat_swiglu_accumulate_token_pair(%tokens, %k, %av, %t4, %kk, %gh, %uh, %acc_g4, %acc_u4) : (index, index, view<[%tokens]x[%k]xf32>, index, index, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) + scf.yield %next_g0, %next_u0, %next_g1, %next_u1, %next_g2, %next_u2, %next_g3, %next_u3, %next_g4, %next_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %fallback_total_g0, %fallback_total_u0, %fallback_total_g1, %fallback_total_u1, %fallback_total_g2, %fallback_total_u2, %fallback_total_g3, %fallback_total_u3, %fallback_total_g4, %fallback_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + scf.yield %branch_total_g0, %branch_total_u0, %branch_total_g1, %branch_total_u1, %branch_total_g2, %branch_total_u2, %branch_total_g3, %branch_total_u3, %branch_total_g4, %branch_total_u4 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> + } + %leader = index.cmp eq, %lane16, %c0 : index + %p0 = index.constant 0 : index + %publishes0 = index.cmp ult, %p0, %tokens : index + scf.if %publishes0 { + %g_lane = vector.reduce %total_g0, %zero_scalar : vector<4xf32>, f32 + %u_lane = vector.reduce %total_u0, %zero_scalar : vector<4xf32>, f32 + %g_width = scalar.constant 32 : i32 + %g_xor1 = scalar.constant 1 : i32 + %g_peer1, %g_ok1 = kernel.subgroup.shuffle %g_lane, %g_xor1, %g_width : f32, i32, i32 + %g_sum1 = scalar.addf %g_lane, %g_peer1 : f32 + %g_xor2 = scalar.constant 2 : i32 + %g_peer2, %g_ok2 = kernel.subgroup.shuffle %g_sum1, %g_xor2, %g_width : f32, i32, i32 + %g_sum2 = scalar.addf %g_sum1, %g_peer2 : f32 + %g_xor4 = scalar.constant 4 : i32 + %g_peer4, %g_ok4 = kernel.subgroup.shuffle %g_sum2, %g_xor4, %g_width : f32, i32, i32 + %g_sum4 = scalar.addf %g_sum2, %g_peer4 : f32 + %g_xor8 = scalar.constant 8 : i32 + %g_peer8, %g_ok8 = kernel.subgroup.shuffle %g_sum4, %g_xor8, %g_width : f32, i32, i32 + %g = scalar.addf %g_sum4, %g_peer8 : f32 + %u_width = scalar.constant 32 : i32 + %u_xor1 = scalar.constant 1 : i32 + %u_peer1, %u_ok1 = kernel.subgroup.shuffle %u_lane, %u_xor1, %u_width : f32, i32, i32 + %u_sum1 = scalar.addf %u_lane, %u_peer1 : f32 + %u_xor2 = scalar.constant 2 : i32 + %u_peer2, %u_ok2 = kernel.subgroup.shuffle %u_sum1, %u_xor2, %u_width : f32, i32, i32 + %u_sum2 = scalar.addf %u_sum1, %u_peer2 : f32 + %u_xor4 = scalar.constant 4 : i32 + %u_peer4, %u_ok4 = kernel.subgroup.shuffle %u_sum2, %u_xor4, %u_width : f32, i32, i32 + %u_sum4 = scalar.addf %u_sum2, %u_peer4 : f32 + %u_xor8 = scalar.constant 8 : i32 + %u_peer8, %u_ok8 = kernel.subgroup.shuffle %u_sum4, %u_xor8, %u_width : f32, i32, i32 + %u = scalar.addf %u_sum4, %u_peer8 : f32 + scf.if %leader { + %gv = vector.splat %g : vector<4xf32> + %uv = vector.splat %u : vector<4xf32> + %result = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gv, %uv) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %value = vector.extract %result[0] : vector<4xf32> -> f32 + view.store %value, %ov[%p0, %row] : f32, view<[%tokens]x[%n]xf32> + } + } + %p1 = index.constant 1 : index + %publishes1 = index.cmp ult, %p1, %tokens : index + scf.if %publishes1 { + %g_lane = vector.reduce %total_g1, %zero_scalar : vector<4xf32>, f32 + %u_lane = vector.reduce %total_u1, %zero_scalar : vector<4xf32>, f32 + %g_width = scalar.constant 32 : i32 + %g_xor1 = scalar.constant 1 : i32 + %g_peer1, %g_ok1 = kernel.subgroup.shuffle %g_lane, %g_xor1, %g_width : f32, i32, i32 + %g_sum1 = scalar.addf %g_lane, %g_peer1 : f32 + %g_xor2 = scalar.constant 2 : i32 + %g_peer2, %g_ok2 = kernel.subgroup.shuffle %g_sum1, %g_xor2, %g_width : f32, i32, i32 + %g_sum2 = scalar.addf %g_sum1, %g_peer2 : f32 + %g_xor4 = scalar.constant 4 : i32 + %g_peer4, %g_ok4 = kernel.subgroup.shuffle %g_sum2, %g_xor4, %g_width : f32, i32, i32 + %g_sum4 = scalar.addf %g_sum2, %g_peer4 : f32 + %g_xor8 = scalar.constant 8 : i32 + %g_peer8, %g_ok8 = kernel.subgroup.shuffle %g_sum4, %g_xor8, %g_width : f32, i32, i32 + %g = scalar.addf %g_sum4, %g_peer8 : f32 + %u_width = scalar.constant 32 : i32 + %u_xor1 = scalar.constant 1 : i32 + %u_peer1, %u_ok1 = kernel.subgroup.shuffle %u_lane, %u_xor1, %u_width : f32, i32, i32 + %u_sum1 = scalar.addf %u_lane, %u_peer1 : f32 + %u_xor2 = scalar.constant 2 : i32 + %u_peer2, %u_ok2 = kernel.subgroup.shuffle %u_sum1, %u_xor2, %u_width : f32, i32, i32 + %u_sum2 = scalar.addf %u_sum1, %u_peer2 : f32 + %u_xor4 = scalar.constant 4 : i32 + %u_peer4, %u_ok4 = kernel.subgroup.shuffle %u_sum2, %u_xor4, %u_width : f32, i32, i32 + %u_sum4 = scalar.addf %u_sum2, %u_peer4 : f32 + %u_xor8 = scalar.constant 8 : i32 + %u_peer8, %u_ok8 = kernel.subgroup.shuffle %u_sum4, %u_xor8, %u_width : f32, i32, i32 + %u = scalar.addf %u_sum4, %u_peer8 : f32 + scf.if %leader { + %gv = vector.splat %g : vector<4xf32> + %uv = vector.splat %u : vector<4xf32> + %result = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gv, %uv) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %value = vector.extract %result[0] : vector<4xf32> -> f32 + view.store %value, %ov[%p1, %row] : f32, view<[%tokens]x[%n]xf32> + } + } + %p2 = index.constant 2 : index + %publishes2 = index.cmp ult, %p2, %tokens : index + scf.if %publishes2 { + %g_lane = vector.reduce %total_g2, %zero_scalar : vector<4xf32>, f32 + %u_lane = vector.reduce %total_u2, %zero_scalar : vector<4xf32>, f32 + %g_width = scalar.constant 32 : i32 + %g_xor1 = scalar.constant 1 : i32 + %g_peer1, %g_ok1 = kernel.subgroup.shuffle %g_lane, %g_xor1, %g_width : f32, i32, i32 + %g_sum1 = scalar.addf %g_lane, %g_peer1 : f32 + %g_xor2 = scalar.constant 2 : i32 + %g_peer2, %g_ok2 = kernel.subgroup.shuffle %g_sum1, %g_xor2, %g_width : f32, i32, i32 + %g_sum2 = scalar.addf %g_sum1, %g_peer2 : f32 + %g_xor4 = scalar.constant 4 : i32 + %g_peer4, %g_ok4 = kernel.subgroup.shuffle %g_sum2, %g_xor4, %g_width : f32, i32, i32 + %g_sum4 = scalar.addf %g_sum2, %g_peer4 : f32 + %g_xor8 = scalar.constant 8 : i32 + %g_peer8, %g_ok8 = kernel.subgroup.shuffle %g_sum4, %g_xor8, %g_width : f32, i32, i32 + %g = scalar.addf %g_sum4, %g_peer8 : f32 + %u_width = scalar.constant 32 : i32 + %u_xor1 = scalar.constant 1 : i32 + %u_peer1, %u_ok1 = kernel.subgroup.shuffle %u_lane, %u_xor1, %u_width : f32, i32, i32 + %u_sum1 = scalar.addf %u_lane, %u_peer1 : f32 + %u_xor2 = scalar.constant 2 : i32 + %u_peer2, %u_ok2 = kernel.subgroup.shuffle %u_sum1, %u_xor2, %u_width : f32, i32, i32 + %u_sum2 = scalar.addf %u_sum1, %u_peer2 : f32 + %u_xor4 = scalar.constant 4 : i32 + %u_peer4, %u_ok4 = kernel.subgroup.shuffle %u_sum2, %u_xor4, %u_width : f32, i32, i32 + %u_sum4 = scalar.addf %u_sum2, %u_peer4 : f32 + %u_xor8 = scalar.constant 8 : i32 + %u_peer8, %u_ok8 = kernel.subgroup.shuffle %u_sum4, %u_xor8, %u_width : f32, i32, i32 + %u = scalar.addf %u_sum4, %u_peer8 : f32 + scf.if %leader { + %gv = vector.splat %g : vector<4xf32> + %uv = vector.splat %u : vector<4xf32> + %result = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gv, %uv) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %value = vector.extract %result[0] : vector<4xf32> -> f32 + view.store %value, %ov[%p2, %row] : f32, view<[%tokens]x[%n]xf32> + } + } + %p3 = index.constant 3 : index + %publishes3 = index.cmp ult, %p3, %tokens : index + scf.if %publishes3 { + %g_lane = vector.reduce %total_g3, %zero_scalar : vector<4xf32>, f32 + %u_lane = vector.reduce %total_u3, %zero_scalar : vector<4xf32>, f32 + %g_width = scalar.constant 32 : i32 + %g_xor1 = scalar.constant 1 : i32 + %g_peer1, %g_ok1 = kernel.subgroup.shuffle %g_lane, %g_xor1, %g_width : f32, i32, i32 + %g_sum1 = scalar.addf %g_lane, %g_peer1 : f32 + %g_xor2 = scalar.constant 2 : i32 + %g_peer2, %g_ok2 = kernel.subgroup.shuffle %g_sum1, %g_xor2, %g_width : f32, i32, i32 + %g_sum2 = scalar.addf %g_sum1, %g_peer2 : f32 + %g_xor4 = scalar.constant 4 : i32 + %g_peer4, %g_ok4 = kernel.subgroup.shuffle %g_sum2, %g_xor4, %g_width : f32, i32, i32 + %g_sum4 = scalar.addf %g_sum2, %g_peer4 : f32 + %g_xor8 = scalar.constant 8 : i32 + %g_peer8, %g_ok8 = kernel.subgroup.shuffle %g_sum4, %g_xor8, %g_width : f32, i32, i32 + %g = scalar.addf %g_sum4, %g_peer8 : f32 + %u_width = scalar.constant 32 : i32 + %u_xor1 = scalar.constant 1 : i32 + %u_peer1, %u_ok1 = kernel.subgroup.shuffle %u_lane, %u_xor1, %u_width : f32, i32, i32 + %u_sum1 = scalar.addf %u_lane, %u_peer1 : f32 + %u_xor2 = scalar.constant 2 : i32 + %u_peer2, %u_ok2 = kernel.subgroup.shuffle %u_sum1, %u_xor2, %u_width : f32, i32, i32 + %u_sum2 = scalar.addf %u_sum1, %u_peer2 : f32 + %u_xor4 = scalar.constant 4 : i32 + %u_peer4, %u_ok4 = kernel.subgroup.shuffle %u_sum2, %u_xor4, %u_width : f32, i32, i32 + %u_sum4 = scalar.addf %u_sum2, %u_peer4 : f32 + %u_xor8 = scalar.constant 8 : i32 + %u_peer8, %u_ok8 = kernel.subgroup.shuffle %u_sum4, %u_xor8, %u_width : f32, i32, i32 + %u = scalar.addf %u_sum4, %u_peer8 : f32 + scf.if %leader { + %gv = vector.splat %g : vector<4xf32> + %uv = vector.splat %u : vector<4xf32> + %result = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gv, %uv) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %value = vector.extract %result[0] : vector<4xf32> -> f32 + view.store %value, %ov[%p3, %row] : f32, view<[%tokens]x[%n]xf32> + } + } + %p4 = index.constant 4 : index + %publishes4 = index.cmp ult, %p4, %tokens : index + scf.if %publishes4 { + %g_lane = vector.reduce %total_g4, %zero_scalar : vector<4xf32>, f32 + %u_lane = vector.reduce %total_u4, %zero_scalar : vector<4xf32>, f32 + %g_width = scalar.constant 32 : i32 + %g_xor1 = scalar.constant 1 : i32 + %g_peer1, %g_ok1 = kernel.subgroup.shuffle %g_lane, %g_xor1, %g_width : f32, i32, i32 + %g_sum1 = scalar.addf %g_lane, %g_peer1 : f32 + %g_xor2 = scalar.constant 2 : i32 + %g_peer2, %g_ok2 = kernel.subgroup.shuffle %g_sum1, %g_xor2, %g_width : f32, i32, i32 + %g_sum2 = scalar.addf %g_sum1, %g_peer2 : f32 + %g_xor4 = scalar.constant 4 : i32 + %g_peer4, %g_ok4 = kernel.subgroup.shuffle %g_sum2, %g_xor4, %g_width : f32, i32, i32 + %g_sum4 = scalar.addf %g_sum2, %g_peer4 : f32 + %g_xor8 = scalar.constant 8 : i32 + %g_peer8, %g_ok8 = kernel.subgroup.shuffle %g_sum4, %g_xor8, %g_width : f32, i32, i32 + %g = scalar.addf %g_sum4, %g_peer8 : f32 + %u_width = scalar.constant 32 : i32 + %u_xor1 = scalar.constant 1 : i32 + %u_peer1, %u_ok1 = kernel.subgroup.shuffle %u_lane, %u_xor1, %u_width : f32, i32, i32 + %u_sum1 = scalar.addf %u_lane, %u_peer1 : f32 + %u_xor2 = scalar.constant 2 : i32 + %u_peer2, %u_ok2 = kernel.subgroup.shuffle %u_sum1, %u_xor2, %u_width : f32, i32, i32 + %u_sum2 = scalar.addf %u_sum1, %u_peer2 : f32 + %u_xor4 = scalar.constant 4 : i32 + %u_peer4, %u_ok4 = kernel.subgroup.shuffle %u_sum2, %u_xor4, %u_width : f32, i32, i32 + %u_sum4 = scalar.addf %u_sum2, %u_peer4 : f32 + %u_xor8 = scalar.constant 8 : i32 + %u_peer8, %u_ok8 = kernel.subgroup.shuffle %u_sum4, %u_xor8, %u_width : f32, i32, i32 + %u = scalar.addf %u_sum4, %u_peer8 : f32 + scf.if %leader { + %gv = vector.splat %g : vector<4xf32> + %uv = vector.splat %u : vector<4xf32> + %result = template.apply<@ggml.mul_mat_swiglu.apply_vector4>(%op, %gv, %uv) : (index, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) + %value = vector.extract %result[0] : vector<4xf32> -> f32 + view.store %value, %ov[%p4, %row] : f32, view<[%tokens]x[%n]xf32> + } + } + } + func.return +} + +kernel.def target(@ggml_mul_mat_swiglu_gfx11_wave32) @ggml_mul_mat_swiglu_f32_f32_lowtoken_dot(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %c1 = index.constant 1 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c128 = index.constant 128 : index + %last_row = index.add %output_size, %c7 : index + %channel_tiles = index.div %last_row, %c8 : index + kernel.launch.config workgroups(%channel_tiles, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 5)] { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %gate_format = config.get @ggml.mul_mat_swiglu.gate_weight_format : index + %up_format = config.get @ggml.mul_mat_swiglu.up_weight_format : index + %tokens = index.assume %token_count [range(%token_count, 1, 5)] : index + %channel_tile = kernel.workgroup.id : index + %op = config.get @ggml.mul_mat_swiglu.op : index + func.call @ggml_mul_mat_swiglu_lowtoken_dot(%op, %gate_format, %up_format, %tokens, %input_size, %output_size, %channel_tile, %input, %gate_weight, %up_weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} + +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +func.def inline @ggml_mul_mat_swiglu_apply_scalar(%op: index, %lhs: f32, %rhs: f32) -> (f32) { + %op_swiglu = index.constant 4 : index + %is_swiglu = index.cmp eq, %op, %op_swiglu : index + %result = scf.if %is_swiglu -> (f32) { + %sigmoid = scalar.logisticf %lhs : f32 + %silu = scalar.mulf %lhs, %sigmoid : f32 + %value = scalar.mulf %silu, %rhs : f32 + scf.yield %value : f32 + } else { + %value = template.apply<@ggml.binary_f32.apply>(%op, %lhs, %rhs) : (index, f32, f32) -> (f32) + scf.yield %value : f32 + } + func.return %result : f32 +} + +// Q4_K weights retain their native affine values; only activations use the +// existing Q8_1_x4 producer. Four 16-lane cohorts share each wave's K256 step. +func.def inline @ggml_swiglu_q8_values(%local: i1, %input: buffer, %stage: buffer, %record_base: offset, %word0: index, %word1: index, %meta0: index, %meta1: index, %sum0: index, %sum1: index) -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) { + %payload_offset = index.constant 16 : offset + %payload_base = index.add %record_base, %payload_offset : offset + %a0, %a1, %d0, %d1, %s0, %s1 = scf.if %local -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) { + %payload = buffer.view %stage[%payload_base] : buffer -> view<32xi32> + %metadata = buffer.view %stage[%record_base] : buffer -> view<8xf16> + %a0 = vector.load %payload[%word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %payload[%word1] : view<32xi32> -> vector<2xi32> + %packet = vector.load %metadata[%meta0] : view<8xf16> -> vector<4xf16> + %values = vector.extf %packet : vector<4xf16> to vector<4xf32> + %d0 = vector.extract %values[0] : vector<4xf32> -> f32 + %d1 = vector.extract %values[2] : vector<4xf32> -> f32 + %s0 = vector.extract %values[1] : vector<4xf32> -> f32 + %s1 = vector.extract %values[3] : vector<4xf32> -> f32 + scf.yield %a0, %a1, %d0, %d1, %s0, %s1 : vector<2xi32>, vector<2xi32>, f32, f32, f32, f32 + } else { + %payload = buffer.view %input[%payload_base] : buffer -> view<32xi32> + %metadata = buffer.view %input[%record_base] : buffer -> view<8xf16> + %a0 = vector.load %payload[%word0] : view<32xi32> -> vector<2xi32> + %a1 = vector.load %payload[%word1] : view<32xi32> -> vector<2xi32> + %d0h = view.load %metadata[%meta0] : view<8xf16> -> f16 + %d1h = view.load %metadata[%meta1] : view<8xf16> -> f16 + %d0 = scalar.extf %d0h : f16 to f32 + %d1 = scalar.extf %d1h : f16 to f32 + %s0h = view.load %metadata[%sum0] : view<8xf16> -> f16 + %s1h = view.load %metadata[%sum1] : view<8xf16> -> f16 + %s0 = scalar.extf %s0h : f16 to f32 + %s1 = scalar.extf %s1h : f16 to f32 + scf.yield %a0, %a1, %d0, %d1, %s0, %s1 : vector<2xi32>, vector<2xi32>, f32, f32, f32, f32 + } + func.return %a0, %a1, %d0, %d1, %s0, %s1 : vector<2xi32>, vector<2xi32>, f32, f32, f32, f32 +} + +func.def inline @ggml_q4_paired_prefetch_words(%gate: buffer, %up: buffer, %group_block: index, %row_offset: offset, %row_lane: index, %payload_field: index, %payload_word: index) -> (vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32>) { + %c0 = index.constant 0 : index + %record_bytes = index.constant 9216 : offset + %payload_add = index.constant 1024 : offset + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %header_base = index.add %group_base, %row_offset : offset + %payload_base = index.add %group_base, %payload_add : offset + %ghv = buffer.view %gate[%header_base] : buffer -> view<4xi32> + %uhv = buffer.view %up[%header_base] : buffer -> view<4xi32> + %gpv = buffer.view %gate[%payload_base] : buffer -> view<8x64x4xi32> + %upv = buffer.view %up[%payload_base] : buffer -> view<8x64x4xi32> + %gh = vector.load %ghv[%c0] : view<4xi32> -> vector<4xi32> + %gr = vector.load %gpv[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + %uh = vector.load %uhv[%c0] : view<4xi32> -> vector<4xi32> + %ur = vector.load %upv[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + func.return %gh, %uh, %gr, %ur : vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32> +} + +// Two K teams amortize native weight service on large paired projections. +// Keep smaller/cache-resident and odd-block reductions on the original mapping. +func.def pure inline @ggml_swiglu_q8_reduction_partitions(%local: i1, %input_size: index, %output_size: index) -> (index) { + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c256 = index.constant 256 : index + %minimum_blocks = index.constant 262144 : index + %blocks = index.div %input_size, %c256 : index + %tail = index.rem %blocks, %c2 : index + %even = index.cmp eq, %tail, %c0 : index + %weight_blocks = index.mul %blocks, %output_size : index + %large = index.cmp uge, %weight_blocks, %minimum_blocks : index + %eligible = scalar.andi %local, %even : i1 + %partitioned = scalar.andi %eligible, %large : i1 + %partitions = scf.select %partitioned, %c2, %c1 : index + func.return %partitions : index +} + +func.decl @ggml_lowtoken_reduce_cohort_f32(%value: f32) -> (f32) + +func.def inline @ggml_swiglu_q8_reduce(%value: f32, %partitioned: i1) -> (f32) { + %width = scalar.constant 32 : i32 + %sum8 = func.call @ggml_lowtoken_reduce_cohort_f32(%value) : (f32) -> (f32) + %result = scf.if %partitioned -> (f32) { + %off16 = scalar.constant 16 : i32 + %peer16, %ok16 = kernel.subgroup.shuffle %sum8, %off16, %width : f32, i32, i32 + %sum = scalar.addf %sum8, %peer16 : f32 + scf.yield %sum : f32 + } else { + %zero = scalar.constant 0.0 : f32 + %sum = scalar.addf %sum8, %zero : f32 + scf.yield %sum : f32 + } + func.return %result : f32 +} + +template.decl @ggml.mul_mat_swiglu.q4_q8_lowtoken_body(%publish_q8: i1, %token_count: index, %input: buffer, %gate: buffer, %up: buffer, %output: 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) + +kernel.def target(@ggml_mul_mat_swiglu_gfx11_wave64) @ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot(%token_count: index) { + %n = config.get @ggml.mul_mat_swiglu.output_size : index + %k = config.get @ggml.mul_mat_swiglu.input_size : index + %capacity = config.get @ggml.workload.token_capacity : index + %local_q8_c1 = index.constant 1 : index + %local_q8_c31 = index.constant 31 : index + %local_q8_c32 = index.constant 32 : index + %local_q8_c128 = index.constant 128 : index + %local_q8_c144 = index.constant 144 : index + %local_q8_c32768 = index.constant 32768 : index + %local_q8_groups = index.div %k, %local_q8_c128 : index + %local_q8_records = index.mul %local_q8_groups, %capacity : index + %local_q8_bytes = index.mul %local_q8_records, %local_q8_c144 : index + %local_q8_padded_rows = index.add %n, %local_q8_c31 : index + %local_q8_workgroups = index.div %local_q8_padded_rows, %local_q8_c32 : index + %local_q8_batched = index.cmp ugt, %capacity, %local_q8_c1 : index + %local_q8_fits = index.cmp ule, %local_q8_bytes, %local_q8_c32768 : index + %local_q8_enough_work = index.cmp uge, %local_q8_workgroups, %local_q8_c128 : index + %local_q8_eligible = scalar.andi %local_q8_batched, %local_q8_fits : i1 + %local = scalar.andi %local_q8_eligible, %local_q8_enough_work : i1 + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c512 = index.constant 512 : index + %partitions = func.call pure @ggml_swiglu_q8_reduction_partitions(%local, %k, %n) : (i1, index, index) -> (index) + %local_rows = index.div %c32, %partitions : index + %rows = scf.select %local, %local_rows, %c8 : index + %padding = index.sub %rows, %c1 : index + %padded_n = index.add %n, %padding : index + %groups = index.div %padded_n, %rows : index + %threads = scf.select %local, %c512, %c128 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%threads, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate: buffer, %up: buffer, %output: buffer) { + %publish_q8 = scalar.constant false : i1 + template.apply<@ggml.mul_mat_swiglu.q4_q8_lowtoken_body>(%publish_q8, %token_count, %input, %gate, %up, %output, %output) : (i1, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_mul_mat_swiglu_gfx11_wave64) @ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output(%token_count: index) { + %n0 = config.get @ggml.mul_mat_swiglu.output_size : index + %n = index.assume %n0 [range(%n0, 1, 262144), mul(%n0, 128)] : index + %k0 = config.get @ggml.mul_mat_swiglu.input_size : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %capacity0 = config.get @ggml.workload.token_capacity : index + %capacity = index.assume %capacity0 [range(%capacity0, 1, 5)] : index + %local_q8_c1 = index.constant 1 : index + %local_q8_c31 = index.constant 31 : index + %local_q8_c32 = index.constant 32 : index + %local_q8_c128 = index.constant 128 : index + %local_q8_c144 = index.constant 144 : index + %local_q8_c32768 = index.constant 32768 : index + %local_q8_groups = index.div %k, %local_q8_c128 : index + %local_q8_records = index.mul %local_q8_groups, %capacity : index + %local_q8_bytes = index.mul %local_q8_records, %local_q8_c144 : index + %local_q8_padded_rows = index.add %n, %local_q8_c31 : index + %local_q8_workgroups = index.div %local_q8_padded_rows, %local_q8_c32 : index + %local_q8_batched = index.cmp ugt, %capacity, %local_q8_c1 : index + %local_q8_fits = index.cmp ule, %local_q8_bytes, %local_q8_c32768 : index + %local_q8_enough_work = index.cmp uge, %local_q8_workgroups, %local_q8_c128 : index + %local_q8_eligible = scalar.andi %local_q8_batched, %local_q8_fits : i1 + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c512 = index.constant 512 : index + %c1024 = index.constant 1024 : index + %local = scalar.andi %local_q8_eligible, %local_q8_enough_work : i1 + %partitions = func.call pure @ggml_swiglu_q8_reduction_partitions(%local, %k, %n) : (i1, index, index) -> (index) + %partitioned = index.cmp ugt, %partitions, %c1 : index + %rows = index.constant 32 : index + %padding = index.sub %rows, %c1 : index + %padded_n = index.add %n, %padding : index + %groups = index.div %padded_n, %rows : index + %threads = scf.select %partitioned, %c1024, %c512 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%threads, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %gate: buffer, %up: buffer, %output: buffer, %q8_output: buffer) where [range(%token_count, 1, 5)] { + %publish_q8 = scalar.constant true : i1 + template.apply<@ggml.mul_mat_swiglu.q4_q8_lowtoken_body>(%publish_q8, %token_count, %input, %gate, %up, %output, %q8_output) : (i1, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +template.def<@ggml.mul_mat_swiglu.q4_q8_lowtoken_body> device @ggml_mul_mat_swiglu_q4_q8_lowtoken_body(%publish_q8: i1, %token_count: index, %input: buffer, %gate: buffer, %up: buffer, %output: buffer, %q8_output: buffer) { + %op = config.get @ggml.mul_mat_swiglu.op : index + %k0 = config.get @ggml.mul_mat_swiglu.input_size : index + %n0 = config.get @ggml.mul_mat_swiglu.output_size : index + %capacity = config.get @ggml.workload.token_capacity : index + %tokens = index.assume %token_count [range(%token_count, 1, 5), le(%token_count, %capacity)] : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %n = index.assume %n0 [range(%n0, 1, 262144)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c128 = index.constant 128 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %block_bytes = index.constant 144 : offset + %code_offset = index.constant 16 : offset + %zero = scalar.constant 0.0 : f32 + %lane0 = kernel.subgroup.lane.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index + %wave = kernel.subgroup.id : index + %wg = kernel.workgroup.id : index + %local_q8_c1 = index.constant 1 : index + %local_q8_c31 = index.constant 31 : index + %local_q8_c32 = index.constant 32 : index + %local_q8_c128 = index.constant 128 : index + %local_q8_c144 = index.constant 144 : index + %local_q8_c32768 = index.constant 32768 : index + %local_q8_groups = index.div %k, %local_q8_c128 : index + %local_q8_records = index.mul %local_q8_groups, %capacity : index + %local_q8_bytes = index.mul %local_q8_records, %local_q8_c144 : index + %local_q8_padded_rows = index.add %n, %local_q8_c31 : index + %local_q8_workgroups = index.div %local_q8_padded_rows, %local_q8_c32 : index + %local_q8_batched = index.cmp ugt, %capacity, %local_q8_c1 : index + %local_q8_fits = index.cmp ule, %local_q8_bytes, %local_q8_c32768 : index + %local_q8_enough_work = index.cmp uge, %local_q8_workgroups, %local_q8_c128 : index + %local_q8_eligible = scalar.andi %local_q8_batched, %local_q8_fits : i1 + %use_local = scalar.andi %local_q8_eligible, %local_q8_enough_work : i1 + %partitions = func.call pure @ggml_swiglu_q8_reduction_partitions(%use_local, %k, %n) : (i1, index, index) -> (index) + %partitioned = index.cmp ugt, %partitions, %c1 : index + %publication_rows = index.constant 32 : index + %publication_row_budget = index.mul %publication_rows, %partitions : index + %local_row_budget = scf.select %publish_q8, %publication_row_budget, %c32 : index + %local_rows = index.div %local_row_budget, %partitions : index + %default_rows_per_workgroup = scf.select %use_local, %local_rows, %c8 : index + %rows_per_workgroup = scf.select %publish_q8, %publication_rows, %default_rows_per_workgroup : index + %cohort_width = index.mul %c16, %partitions : index + %wave_rows = index.div %c4, %partitions : index + %row_base = index.mul %wg, %rows_per_workgroup : index + %wave_row = index.mul %wave, %wave_rows : index + %row_half = index.div %lane, %cohort_width : index + %local_row = index.add %wave_row, %row_half : index + %row = index.add %row_base, %local_row : index + %valid_row = index.cmp ult, %row, %n : index + %cohort = index.div %lane, %cohort_width : index + %team_lane = index.rem %lane, %cohort_width : index + %partition = index.div %team_lane, %c16 : index + %block_lane = index.rem %lane, %c16 : index + %pair = index.div %block_lane, %c4 : index + %packet = index.rem %block_lane, %c4 : index + %group0 = index.mul %pair, %c2 : index + %group1 = index.add %group0, %c1 : index + %word_page = index.mul %pair, %c8 : index + %word_add = index.mul %packet, %c2 : index + %word_index = index.add %word_page, %word_add : index + %meta_lane = index.cmp eq, %packet, %c0 : index + %blocks = index.div %k, %c256 : index + %partition_blocks = index.div %blocks, %partitions : index + %partition_origin = index.mul %partition, %partition_blocks : index + %weight_row_bytes = index.scale %blocks, %block_bytes : index, offset -> offset + %weight_row_base = index.scale %row, %weight_row_bytes : index, offset -> offset + %q8_groups = index.div %k, %c128 : index + %q8_row_bytes = index.scale %q8_groups, %block_bytes : index, offset -> offset + %q8_half = index.div %pair, %c2 : index + %q8_pair = index.rem %pair, %c2 : index + %q8_group0 = index.mul %q8_pair, %c2 : index + %q8_group1 = index.add %q8_group0, %c1 : index + %q8_words_base = index.mul %q8_group0, %c8 : index + %q8_word0 = index.add %q8_words_base, %word_add : index + %q8_word1 = index.add %q8_word0, %c8 : index + %q8_meta0 = index.mul %q8_group0, %c2 : index + %q8_meta1 = index.mul %q8_group1, %c2 : index + %q8_sum0 = index.add %q8_meta0, %c1 : index + %q8_sum1 = index.add %q8_meta1, %c1 : index + %c64 = index.constant 64 : index + %record_bytes = index.constant 9216 : offset + %row_field_bytes = index.constant 16 : offset + %payload_add = index.constant 1024 : offset + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %row_group_start = index.mul %row_group, %blocks : index + %row_group_block = index.add %row_group_start, %partition_origin : index + %row_offset = index.scale %row_lane, %row_field_bytes : index, offset -> offset + %payload_field = index.div %word_index, %c4 : index + %payload_word = index.rem %word_index, %c4 : index + %input_na, %gate_na, %up_na, %output_na = buffer.assume.noalias %input, %gate, %up, %output : buffer, buffer, buffer, buffer + // One immutable Q8 activation record stream is shared by all row cohorts. + %activation_rows = index.mul %capacity, %q8_groups : index + %activation_active_rows = index.mul %tokens, %q8_groups : index + %activation_bytes = index.scale %activation_rows, %block_bytes : index, offset -> offset + %local_bytes = scf.select %use_local, %activation_bytes, %base : offset + %activation_stage = buffer.alloca align(16) %local_bytes : buffer + %epilogue_capacity_bytes = index.constant 640 : offset + %epilogue_bytes = scf.select %publish_q8, %epilogue_capacity_bytes, %base : offset + %epilogue_stage = buffer.alloca align(16) %epilogue_bytes : buffer + %epilogue_view = buffer.view %epilogue_stage[%base] : buffer -> view<5x32xf32> + %epilogue_scratch_capacity_bytes = index.constant 1024 : offset + %epilogue_scratch_bytes = scf.select %publish_q8, %epilogue_scratch_capacity_bytes, %base : offset + %epilogue_scratch_stage = buffer.alloca align(16) %epilogue_scratch_bytes : buffer + %epilogue_row = index.rem %row, %c32 : index + scf.if %use_local { + %record_words = index.constant 36 : index + %record_packets = index.constant 9 : index + %activation_words = index.mul %activation_rows, %record_words : index + %activation_packets = index.mul %activation_rows, %record_packets : index + %active_packets = index.mul %activation_active_rows, %record_packets : index + %source_words = buffer.view %input_na[%base] : buffer -> view<[%activation_words]xi32> + %stage_words = buffer.view %activation_stage[%base] : buffer -> view<[%activation_words]xi32> + %stage_workgroup = kernel.workgroup.size : index + %stage_round_add = index.sub %stage_workgroup, %c1 : index + %stage_rounded = index.add %activation_packets, %stage_round_add : index + %stage_rounds = index.div %stage_rounded, %stage_workgroup : index + %stage_thread = kernel.workitem.id : index + scf.for %round = [%c0 to %stage_rounds step %c1] { + %round_origin = index.mul %round, %stage_workgroup : index + %packet = index.add %round_origin, %stage_thread : index + %live_packet = index.cmp ult, %packet, %active_packets : index + scf.if %live_packet { + %word = index.mul %packet, %c4 : index + %words = vector.load %source_words[%word] : view<[%activation_words]xi32> -> vector<4xi32> + vector.store %words, %stage_words[%word] : vector<4xi32>, view<[%activation_words]xi32> + scf.yield + } + scf.yield + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield + } + scf.if %valid_row { + %initial_g_headers, %initial_u_headers, %initial_g_raw, %initial_u_raw = scf.if %use_local -> (vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32>) { + %gh, %uh, %gr, %ur = func.call @ggml_q4_paired_prefetch_words(%gate_na, %up_na, %row_group_block, %row_offset, %row_lane, %payload_field, %payload_word) : (buffer, buffer, index, offset, index, index, index) -> (vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32>) + scf.yield %gh, %uh, %gr, %ur : vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32> + } else { + %zero_header = vector.constant 0 : vector<4xi32> + %zero_payload = vector.constant 0 : vector<2xi32> + scf.yield %zero_header, %zero_header, %zero_payload, %zero_payload : vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32> + } + %total_g0, %total_g1, %total_g2, %total_g3, %total_g4, %total_u0, %total_u1, %total_u2, %total_u3, %total_u4, %unused_g_headers, %unused_u_headers, %unused_g_raw, %unused_u_raw = scf.for %block_base = [%c0 to %partition_blocks step %c1](%acc_g0 = %zero : f32, %acc_g1 = %zero : f32, %acc_g2 = %zero : f32, %acc_g3 = %zero : f32, %acc_g4 = %zero : f32, %acc_u0 = %zero : f32, %acc_u1 = %zero : f32, %acc_u2 = %zero : f32, %acc_u3 = %zero : f32, %acc_u4 = %zero : f32, %current_g_headers = %initial_g_headers : vector<4xi32>, %current_u_headers = %initial_u_headers : vector<4xi32>, %current_g_raw = %initial_g_raw : vector<2xi32>, %current_u_raw = %initial_u_raw : vector<2xi32>) -> (f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32>) { + %block = index.add %block_base, %partition_origin : index + %valid_block = index.cmp ult, %block, %blocks : index + %safe_block = scf.select %valid_block, %block, %c0 : index + %group_block = index.add %row_group_start, %safe_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %wb = index.add %group_base, %row_offset : offset + %wc = index.add %group_base, %payload_add : offset + %g_header_view = buffer.view %gate_na[%wb] : buffer -> view<4xi32> + %g_code_view = buffer.view %gate_na[%wc] : buffer -> view<8x64x4xi32> + %g_headers = scf.if %use_local -> (vector<4xi32>) { + scf.yield %current_g_headers : vector<4xi32> + } else { + %loaded = vector.load %g_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %loaded : vector<4xi32> + } + %g_dm_i32 = vector.extract %g_headers[0] : vector<4xi32> -> i32 + %g_dm_word = vector.splat %g_dm_i32 : vector<1xi32> + %g_dm = vector.bitcast %g_dm_word : vector<1xi32> to vector<2xf16> + %g_d0h = vector.extract %g_dm[0] : vector<2xf16> -> f16 + %g_d1h = vector.extract %g_dm[1] : vector<2xf16> -> f16 + %g_d0 = scalar.extf %g_d0h : f16 to f32 + %g_d1 = scalar.extf %g_d1h : f16 to f32 + %g_s0 = vector.extract %g_headers[1] : vector<4xi32> -> i32 + %g_s1 = vector.extract %g_headers[2] : vector<4xi32> -> i32 + %g_s2 = vector.extract %g_headers[3] : vector<4xi32> -> i32 + %g_scale0i, %g_min0i = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + %g_scale0f = scalar.uitofp %g_scale0i : i32 to f32 + %g_min0f = scalar.uitofp %g_min0i : i32 to f32 + %g_ws0 = scalar.mulf %g_d0, %g_scale0f : f32 + %g_wm0 = scalar.mulf %g_d1, %g_min0f : f32 + %g_scale1i, %g_min1i = func.call @ggml_q4k_scale_min_from_header(%g_s0, %g_s1, %g_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + %g_scale1f = scalar.uitofp %g_scale1i : i32 to f32 + %g_min1f = scalar.uitofp %g_min1i : i32 to f32 + %g_ws1 = scalar.mulf %g_d0, %g_scale1f : f32 + %g_wm1 = scalar.mulf %g_d1, %g_min1f : f32 + %g_raw = scf.if %use_local -> (vector<2xi32>) { + scf.yield %current_g_raw : vector<2xi32> + } else { + %loaded = vector.load %g_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %loaded : vector<2xi32> + } + %g_mask = vector.constant 252645135 : vector<2xi32> + %g_shift = vector.constant 4 : vector<2xi32> + %g_lo = vector.andi %g_raw, %g_mask : vector<2xi32> + %g_high = vector.shrui %g_raw, %g_shift : vector<2xi32> + %g_hi = vector.andi %g_high, %g_mask : vector<2xi32> + %u_header_view = buffer.view %up_na[%wb] : buffer -> view<4xi32> + %u_code_view = buffer.view %up_na[%wc] : buffer -> view<8x64x4xi32> + %u_headers = scf.if %use_local -> (vector<4xi32>) { + scf.yield %current_u_headers : vector<4xi32> + } else { + %loaded = vector.load %u_header_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %loaded : vector<4xi32> + } + %u_dm_i32 = vector.extract %u_headers[0] : vector<4xi32> -> i32 + %u_dm_word = vector.splat %u_dm_i32 : vector<1xi32> + %u_dm = vector.bitcast %u_dm_word : vector<1xi32> to vector<2xf16> + %u_d0h = vector.extract %u_dm[0] : vector<2xf16> -> f16 + %u_d1h = vector.extract %u_dm[1] : vector<2xf16> -> f16 + %u_d0 = scalar.extf %u_d0h : f16 to f32 + %u_d1 = scalar.extf %u_d1h : f16 to f32 + %u_s0 = vector.extract %u_headers[1] : vector<4xi32> -> i32 + %u_s1 = vector.extract %u_headers[2] : vector<4xi32> -> i32 + %u_s2 = vector.extract %u_headers[3] : vector<4xi32> -> i32 + %u_scale0i, %u_min0i = func.call @ggml_q4k_scale_min_from_header(%u_s0, %u_s1, %u_s2, %group0) : (i32, i32, i32, index) -> (i32, i32) + %u_scale0f = scalar.uitofp %u_scale0i : i32 to f32 + %u_min0f = scalar.uitofp %u_min0i : i32 to f32 + %u_ws0 = scalar.mulf %u_d0, %u_scale0f : f32 + %u_wm0 = scalar.mulf %u_d1, %u_min0f : f32 + %u_scale1i, %u_min1i = func.call @ggml_q4k_scale_min_from_header(%u_s0, %u_s1, %u_s2, %group1) : (i32, i32, i32, index) -> (i32, i32) + %u_scale1f = scalar.uitofp %u_scale1i : i32 to f32 + %u_min1f = scalar.uitofp %u_min1i : i32 to f32 + %u_ws1 = scalar.mulf %u_d0, %u_scale1f : f32 + %u_wm1 = scalar.mulf %u_d1, %u_min1f : f32 + %u_raw = scf.if %use_local -> (vector<2xi32>) { + scf.yield %current_u_raw : vector<2xi32> + } else { + %loaded = vector.load %u_code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<2xi32> + scf.yield %loaded : vector<2xi32> + } + %u_mask = vector.constant 252645135 : vector<2xi32> + %u_shift = vector.constant 4 : vector<2xi32> + %u_lo = vector.andi %u_raw, %u_mask : vector<2xi32> + %u_high = vector.shrui %u_raw, %u_shift : vector<2xi32> + %u_hi = vector.andi %u_high, %u_mask : vector<2xi32> + %q8_block0 = index.mul %safe_block, %c2 : index + %q8_block = index.add %q8_block0, %q8_half : index + %q8_block_add = index.scale %q8_block, %block_bytes : index, offset -> offset + scf.if %use_local { + scf.schedule.fence + scf.yield + } + %next_block = index.add %block_base, %c1 : index + %has_next = index.cmp ult, %next_block, %partition_blocks : index + %prefetch = scalar.andi %use_local, %has_next : i1 + %next_g_headers, %next_u_headers, %next_g_raw, %next_u_raw = scf.if %prefetch -> (vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32>) { + %next_group_block = index.add %row_group_block, %next_block : index + %gh, %uh, %gr, %ur = func.call @ggml_q4_paired_prefetch_words(%gate_na, %up_na, %next_group_block, %row_offset, %row_lane, %payload_field, %payload_word) : (buffer, buffer, index, offset, index, index, index) -> (vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32>) + scf.yield %gh, %uh, %gr, %ur : vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32> + } else { + %zero_header = vector.constant 0 : vector<4xi32> + %zero_payload = vector.constant 0 : vector<2xi32> + scf.yield %zero_header, %zero_header, %zero_payload, %zero_payload : vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32> + } + %t0 = index.constant 0 : index + %valid0 = index.cmp ult, %t0, %tokens : index + %next_g0, %next_u0 = scf.if %valid0 -> (f32, f32) { + %tb = index.scale %t0, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %a0, %a1, %ad0, %ad1, %as0_raw, %as1_raw = func.call @ggml_swiglu_q8_values(%use_local, %input_na, %activation_stage, %qb, %q8_word0, %q8_word1, %q8_meta0, %q8_meta1, %q8_sum0, %q8_sum1) : (i1, buffer, buffer, offset, index, index, index, index, index, index) -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%g_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%g_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g0, %g_valid : f32 + %u_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%u_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%u_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_s0 = scalar.mulf %u_ws0, %ad0 : f32 + %u_s1 = scalar.mulf %u_ws1, %ad1 : f32 + %u_c0 = scalar.mulf %u_dot0, %u_s0 : f32 + %u_contribution = scalar.fmaf %u_dot1, %u_s1, %u_c0 : f32 + %u_cor0 = scalar.mulf %u_wm0, %as0 : f32 + %u_cor = scalar.fmaf %u_wm1, %as1, %u_cor0 : f32 + %u_corrected = scalar.subf %u_contribution, %u_cor : f32 + %u_valid = scf.select %valid_block, %u_corrected, %zero : f32 + %u_updated = scalar.addf %acc_u0, %u_valid : f32 + scf.yield %g_updated, %u_updated : f32, f32 + } else { + scf.yield %acc_g0, %acc_u0 : f32, f32 + } + %t1 = index.constant 1 : index + %valid1 = index.cmp ult, %t1, %tokens : index + %next_g1, %next_u1 = scf.if %valid1 -> (f32, f32) { + %tb = index.scale %t1, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %a0, %a1, %ad0, %ad1, %as0_raw, %as1_raw = func.call @ggml_swiglu_q8_values(%use_local, %input_na, %activation_stage, %qb, %q8_word0, %q8_word1, %q8_meta0, %q8_meta1, %q8_sum0, %q8_sum1) : (i1, buffer, buffer, offset, index, index, index, index, index, index) -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%g_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%g_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g1, %g_valid : f32 + %u_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%u_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%u_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_s0 = scalar.mulf %u_ws0, %ad0 : f32 + %u_s1 = scalar.mulf %u_ws1, %ad1 : f32 + %u_c0 = scalar.mulf %u_dot0, %u_s0 : f32 + %u_contribution = scalar.fmaf %u_dot1, %u_s1, %u_c0 : f32 + %u_cor0 = scalar.mulf %u_wm0, %as0 : f32 + %u_cor = scalar.fmaf %u_wm1, %as1, %u_cor0 : f32 + %u_corrected = scalar.subf %u_contribution, %u_cor : f32 + %u_valid = scf.select %valid_block, %u_corrected, %zero : f32 + %u_updated = scalar.addf %acc_u1, %u_valid : f32 + scf.yield %g_updated, %u_updated : f32, f32 + } else { + scf.yield %acc_g1, %acc_u1 : f32, f32 + } + scf.if %publish_q8 { + scf.schedule.fence + } + %t2 = index.constant 2 : index + %valid2 = index.cmp ult, %t2, %tokens : index + %next_g2, %next_u2 = scf.if %valid2 -> (f32, f32) { + %tb = index.scale %t2, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %a0, %a1, %ad0, %ad1, %as0_raw, %as1_raw = func.call @ggml_swiglu_q8_values(%use_local, %input_na, %activation_stage, %qb, %q8_word0, %q8_word1, %q8_meta0, %q8_meta1, %q8_sum0, %q8_sum1) : (i1, buffer, buffer, offset, index, index, index, index, index, index) -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%g_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%g_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g2, %g_valid : f32 + %u_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%u_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%u_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_s0 = scalar.mulf %u_ws0, %ad0 : f32 + %u_s1 = scalar.mulf %u_ws1, %ad1 : f32 + %u_c0 = scalar.mulf %u_dot0, %u_s0 : f32 + %u_contribution = scalar.fmaf %u_dot1, %u_s1, %u_c0 : f32 + %u_cor0 = scalar.mulf %u_wm0, %as0 : f32 + %u_cor = scalar.fmaf %u_wm1, %as1, %u_cor0 : f32 + %u_corrected = scalar.subf %u_contribution, %u_cor : f32 + %u_valid = scf.select %valid_block, %u_corrected, %zero : f32 + %u_updated = scalar.addf %acc_u2, %u_valid : f32 + scf.yield %g_updated, %u_updated : f32, f32 + } else { + scf.yield %acc_g2, %acc_u2 : f32, f32 + } + %t3 = index.constant 3 : index + %valid3 = index.cmp ult, %t3, %tokens : index + %next_g3, %next_u3 = scf.if %valid3 -> (f32, f32) { + %tb = index.scale %t3, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %a0, %a1, %ad0, %ad1, %as0_raw, %as1_raw = func.call @ggml_swiglu_q8_values(%use_local, %input_na, %activation_stage, %qb, %q8_word0, %q8_word1, %q8_meta0, %q8_meta1, %q8_sum0, %q8_sum1) : (i1, buffer, buffer, offset, index, index, index, index, index, index) -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%g_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%g_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g3, %g_valid : f32 + %u_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%u_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%u_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_s0 = scalar.mulf %u_ws0, %ad0 : f32 + %u_s1 = scalar.mulf %u_ws1, %ad1 : f32 + %u_c0 = scalar.mulf %u_dot0, %u_s0 : f32 + %u_contribution = scalar.fmaf %u_dot1, %u_s1, %u_c0 : f32 + %u_cor0 = scalar.mulf %u_wm0, %as0 : f32 + %u_cor = scalar.fmaf %u_wm1, %as1, %u_cor0 : f32 + %u_corrected = scalar.subf %u_contribution, %u_cor : f32 + %u_valid = scf.select %valid_block, %u_corrected, %zero : f32 + %u_updated = scalar.addf %acc_u3, %u_valid : f32 + scf.yield %g_updated, %u_updated : f32, f32 + } else { + scf.yield %acc_g3, %acc_u3 : f32, f32 + } + scf.if %publish_q8 { + scf.schedule.fence + } + %t4 = index.constant 4 : index + %valid4 = index.cmp ult, %t4, %tokens : index + %next_g4, %next_u4 = scf.if %valid4 -> (f32, f32) { + %tb = index.scale %t4, %q8_row_bytes : index, offset -> offset + %qb = index.add %tb, %q8_block_add : offset + %a0, %a1, %ad0, %ad1, %as0_raw, %as1_raw = func.call @ggml_swiglu_q8_values(%use_local, %input_na, %activation_stage, %qb, %q8_word0, %q8_word1, %q8_meta0, %q8_meta1, %q8_sum0, %q8_sum1) : (i1, buffer, buffer, offset, index, index, index, index, index, index) -> (vector<2xi32>, vector<2xi32>, f32, f32, f32, f32) + %as0 = scf.select %meta_lane, %as0_raw, %zero : f32 + %as1 = scf.select %meta_lane, %as1_raw, %zero : f32 + %g_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%g_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%g_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %g_s0 = scalar.mulf %g_ws0, %ad0 : f32 + %g_s1 = scalar.mulf %g_ws1, %ad1 : f32 + %g_c0 = scalar.mulf %g_dot0, %g_s0 : f32 + %g_contribution = scalar.fmaf %g_dot1, %g_s1, %g_c0 : f32 + %g_cor0 = scalar.mulf %g_wm0, %as0 : f32 + %g_cor = scalar.fmaf %g_wm1, %as1, %g_cor0 : f32 + %g_corrected = scalar.subf %g_contribution, %g_cor : f32 + %g_valid = scf.select %valid_block, %g_corrected, %zero : f32 + %g_updated = scalar.addf %acc_g4, %g_valid : f32 + %u_dot0 = func.call @ggml_dot_u8_s8_vector8_f32(%u_lo, %a0) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_dot1 = func.call @ggml_dot_u8_s8_vector8_f32(%u_hi, %a1) : (vector<2xi32>, vector<2xi32>) -> (f32) + %u_s0 = scalar.mulf %u_ws0, %ad0 : f32 + %u_s1 = scalar.mulf %u_ws1, %ad1 : f32 + %u_c0 = scalar.mulf %u_dot0, %u_s0 : f32 + %u_contribution = scalar.fmaf %u_dot1, %u_s1, %u_c0 : f32 + %u_cor0 = scalar.mulf %u_wm0, %as0 : f32 + %u_cor = scalar.fmaf %u_wm1, %as1, %u_cor0 : f32 + %u_corrected = scalar.subf %u_contribution, %u_cor : f32 + %u_valid = scf.select %valid_block, %u_corrected, %zero : f32 + %u_updated = scalar.addf %acc_u4, %u_valid : f32 + scf.yield %g_updated, %u_updated : f32, f32 + } else { + scf.yield %acc_g4, %acc_u4 : f32, f32 + } + scf.yield %next_g0, %next_g1, %next_g2, %next_g3, %next_g4, %next_u0, %next_u1, %next_u2, %next_u3, %next_u4, %next_g_headers, %next_u_headers, %next_g_raw, %next_u_raw : f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, vector<4xi32>, vector<4xi32>, vector<2xi32>, vector<2xi32> + } + %leader = index.cmp eq, %team_lane, %c0 : index + %out_view = buffer.view %output_na[%base] : buffer -> view<[%tokens]x[%n]xf32> + %p0 = index.constant 0 : index + %pv0 = index.cmp ult, %p0, %tokens : index + scf.if %pv0 { + %g_sum = func.call @ggml_swiglu_q8_reduce(%total_g0, %partitioned) : (f32, i1) -> (f32) + %u_sum = func.call @ggml_swiglu_q8_reduce(%total_u0, %partitioned) : (f32, i1) -> (f32) + scf.if %leader { + %value = func.call @ggml_mul_mat_swiglu_apply_scalar(%op, %g_sum, %u_sum) : (index, f32, f32) -> (f32) + view.store %value, %out_view[%p0, %row] : f32, view<[%tokens]x[%n]xf32> + scf.if %publish_q8 { + view.store %value, %epilogue_view[%p0, %epilogue_row] : f32, view<5x32xf32> + } + } + } + %p1 = index.constant 1 : index + %pv1 = index.cmp ult, %p1, %tokens : index + scf.if %pv1 { + %g_sum = func.call @ggml_swiglu_q8_reduce(%total_g1, %partitioned) : (f32, i1) -> (f32) + %u_sum = func.call @ggml_swiglu_q8_reduce(%total_u1, %partitioned) : (f32, i1) -> (f32) + scf.if %leader { + %value = func.call @ggml_mul_mat_swiglu_apply_scalar(%op, %g_sum, %u_sum) : (index, f32, f32) -> (f32) + view.store %value, %out_view[%p1, %row] : f32, view<[%tokens]x[%n]xf32> + scf.if %publish_q8 { + view.store %value, %epilogue_view[%p1, %epilogue_row] : f32, view<5x32xf32> + } + } + } + %p2 = index.constant 2 : index + %pv2 = index.cmp ult, %p2, %tokens : index + scf.if %pv2 { + %g_sum = func.call @ggml_swiglu_q8_reduce(%total_g2, %partitioned) : (f32, i1) -> (f32) + %u_sum = func.call @ggml_swiglu_q8_reduce(%total_u2, %partitioned) : (f32, i1) -> (f32) + scf.if %leader { + %value = func.call @ggml_mul_mat_swiglu_apply_scalar(%op, %g_sum, %u_sum) : (index, f32, f32) -> (f32) + view.store %value, %out_view[%p2, %row] : f32, view<[%tokens]x[%n]xf32> + scf.if %publish_q8 { + view.store %value, %epilogue_view[%p2, %epilogue_row] : f32, view<5x32xf32> + } + } + } + %p3 = index.constant 3 : index + %pv3 = index.cmp ult, %p3, %tokens : index + scf.if %pv3 { + %g_sum = func.call @ggml_swiglu_q8_reduce(%total_g3, %partitioned) : (f32, i1) -> (f32) + %u_sum = func.call @ggml_swiglu_q8_reduce(%total_u3, %partitioned) : (f32, i1) -> (f32) + scf.if %leader { + %value = func.call @ggml_mul_mat_swiglu_apply_scalar(%op, %g_sum, %u_sum) : (index, f32, f32) -> (f32) + view.store %value, %out_view[%p3, %row] : f32, view<[%tokens]x[%n]xf32> + scf.if %publish_q8 { + view.store %value, %epilogue_view[%p3, %epilogue_row] : f32, view<5x32xf32> + } + } + } + %p4 = index.constant 4 : index + %pv4 = index.cmp ult, %p4, %tokens : index + scf.if %pv4 { + %g_sum = func.call @ggml_swiglu_q8_reduce(%total_g4, %partitioned) : (f32, i1) -> (f32) + %u_sum = func.call @ggml_swiglu_q8_reduce(%total_u4, %partitioned) : (f32, i1) -> (f32) + scf.if %leader { + %value = func.call @ggml_mul_mat_swiglu_apply_scalar(%op, %g_sum, %u_sum) : (index, f32, f32) -> (f32) + view.store %value, %out_view[%p4, %row] : f32, view<[%tokens]x[%n]xf32> + scf.if %publish_q8 { + view.store %value, %epilogue_view[%p4, %epilogue_row] : f32, view<5x32xf32> + } + } + } + } + scf.if %publish_q8 { + kernel.barrier scope(workgroup) ordering(acq_rel) + %epilogue_tid = kernel.workitem.id : index + %epilogue_threads = index.mul %tokens, %c8 : index + %epilogue_active = index.cmp ult, %epilogue_tid, %epilogue_threads : index + scf.if %epilogue_active { + %epilogue_token = index.div %epilogue_tid, %c8 : index + %epilogue_word = index.rem %epilogue_tid, %c8 : index + %epilogue_col = index.mul %epilogue_word, %c4 : index + %epilogue_linear = index.madd %epilogue_token, %c32, %epilogue_col : index + %epilogue_flat = buffer.view %epilogue_stage[%base] : buffer -> view<160xf32> + %epilogue_values = vector.load %epilogue_flat[%epilogue_linear] : view<160xf32> -> vector<4xf32> + %epilogue_channel = index.madd %wg, %c32, %epilogue_col : index + %epilogue_groups = index.div %n, %c128 : index + %epilogue_token_groups = index.mul %epilogue_token, %epilogue_groups : index + %epilogue_token_bytes = index.scale %epilogue_token_groups, %block_bytes : index, offset -> offset + %epilogue_scratch = buffer.view %epilogue_scratch_stage[%base] : buffer -> view<256xf32> + %epilogue_scratch_d = buffer.view %epilogue_scratch_stage[%base] : buffer -> view<32xf32> + %epilogue_publish = scalar.constant true : i1 + template.apply<@ggml.quantize_q8_1_x4.publish_vector4_strict>(%epilogue_publish, %epilogue_token_bytes, %epilogue_channel, %epilogue_values, %epilogue_scratch, %epilogue_scratch_d, %q8_output) : (i1, offset, index, vector<4xf32>, view<256xf32>, view<32xf32>, buffer) + } + } + template.return +} + +func.decl @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size: index, %output_size: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer) + +// Enough complete M512 tiles and workgroups to amortize the private F16 copy. +func.def pure inline @ggml_swiglu_use_packed_f16(%tokens: index, %outputs: index) -> (i1) { + %zero = index.constant 0 : index + %c64 = index.constant 64 : index + %c512 = index.constant 512 : index + %tail = index.rem %tokens, %c512 : index + %complete = index.cmp eq, %tail, %zero : index + %m_tiles = index.div %tokens, %c512 : index + %n_tiles = index.div %outputs, %c64 : index + %groups = index.mul %m_tiles, %n_tiles : index + %enough_groups = index.cmp uge, %groups, %c64 : index + %packed = scalar.andi %complete, %enough_groups : i1 + func.return %packed : i1 +} + +kernel.def target(@ggml_mul_mat_swiglu_gfx11_wave32) export("ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32") @ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32(%token_count: index) { + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %c1 = index.constant 1 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %c512 = index.constant 512 : index + %c32 = index.constant 32 : index + %packed_input = func.call pure @ggml_swiglu_use_packed_f16(%token_capacity, %output_size) : (index, index) -> (i1) + %token_tile = scf.select %packed_input, %c512, %c256 : index + %channel_tile = scf.select %packed_input, %c32, %c64 : index + %output_tiles = index.div %output_size, %channel_tile : index + %token_tiles = index.div %token_capacity, %token_tile : index + kernel.launch.config workgroups(%token_tiles, %output_tiles, %c1) workgroup_size(%c512, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer) { + %input_size = config.get @ggml.mul_mat_swiglu.input_size : index + %output_size = config.get @ggml.mul_mat_swiglu.output_size : index + %token_capacity = config.get @ggml.workload.token_capacity : index + %packed_input = func.call pure @ggml_swiglu_use_packed_f16(%token_capacity, %output_size) : (index, index) -> (i1) + %bounded_token_count = scf.if %packed_input -> (index) { + // The dispatch specializes and launches the same complete input extent. + %count, %capacity = index.assume %token_count, %token_capacity [range(%token_count, 512, 2048), mul(%token_count, 512), eq(%token_count, %token_capacity)] : index, index + scf.yield %count : index + } else { + %count = index.assume %token_count [range(%token_count, 256, 2048), mul(%token_count, 256), le(%token_count, %token_capacity)] : index + scf.yield %count : index + } + %bounded_output_size = index.assume %output_size [range(%output_size, 64, 262144), mul(%output_size, 64)] : index + %channel_tile = kernel.workgroup.id : index + %token_tile = kernel.workgroup.id : index + %paired = scalar.constant true : i1 + %prefill_weight_format = index.constant 44 : index + %binary_op = config.get @ggml.mul_mat_swiglu.op : index + %f16_output_layout = config.get @ggml.mul_mat_swiglu.f16_output_layout : index + %packed_layout = index.constant 1 : index + %packed_output = index.cmp eq, %f16_output_layout, %packed_layout : index + func.call @ggml_mul_mat_quantized_f16_prefill_wave32(%prefill_weight_format, %binary_op, %paired, %packed_input, %packed_output, %bounded_token_count, %input_size, %bounded_output_size, %channel_tile, %token_tile, %input, %weight, %up_weight, %output, %f16_output) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_symmetric_i4_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_symmetric_i4_wmma.loom new file mode 100644 index 000000000000..daddfb757b9b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_swiglu_symmetric_i4_wmma.loom @@ -0,0 +1,1646 @@ +// Generic dense gate/up contraction using backend-resident symmetric-I4/K32 +// weights and activations. The fused kernel applies SwiGLU and publishes a +// Q8_1 plane consumed by a following projection in the same graph match. +// Shape and source-quant selection belong to the graph dispatcher; this module +// contains no model identity or fixed model dimensions. + +template.decl @ggml.quantize_symmetric_i4_k32.body(%plane_major: i1, %src: buffer, %qs: buffer, %ds: buffer, %sums: buffer) + +// Sixteen lanes quantize one K64 block; each packs four values while the two K32 halves share a scale and keep separate sums. +config.decl @ggml.quantize_symmetric_i4_k32.input_size : %value: index where [range(%value, 32, 65536), mul(%value, 32)] + +config.decl @ggml.quantize_symmetric_i4_k32.token_count : %value: index where [range(%value, 1, 16384)] + +kernel.def export("ggml_quantize_f32_symmetric_i4_k32") @ggml_quantize_f32_symmetric_i4_k32() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %wg_groups_m1 = index.constant 127 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %k = config.get @ggml.quantize_symmetric_i4_k32.input_size : index + %cols = config.get @ggml.quantize_symmetric_i4_k32.token_count : index + %total = index.mul %k, %cols : index + %groups = index.div %total, %c64 : index + %rounded = index.add %groups, %wg_groups_m1 : index + %workgroups = index.div %rounded, %c128 : index + kernel.launch.config workgroups(%workgroups, %unit, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%src: buffer, %qs: buffer, %ds: buffer, %sums: buffer) { + %plane_major = scalar.constant false : i1 + template.apply<@ggml.quantize_symmetric_i4_k32.body>(%plane_major, %src, %qs, %ds, %sums) : (i1, buffer, buffer, buffer, buffer) + kernel.return +} + +template.def<@ggml.quantize_symmetric_i4_k32.body> device @ggml_quantize_f32_symmetric_i4_k32_body(%plane_major: i1, %src: buffer, %qs: buffer, %ds: buffer, %sums: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %amax_eps = scalar.constant 1.0000000000000001e-30 : f32 + %inv7 = scalar.constant 0.14285714285714285 : f32 + %f7 = scalar.constant 7.0 : f32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %shift4 = scalar.constant 4 : i32 + %shift8 = scalar.constant 8 : i32 + %shift12 = scalar.constant 12 : i32 + %shuffle_width = scalar.constant 32 : i32 + %mask15 = vector.constant 15 : vector<4xi32> + + %k = config.get @ggml.quantize_symmetric_i4_k32.input_size : index + %cols = config.get @ggml.quantize_symmetric_i4_k32.token_count : index + %k_b = index.assume %k [range(%k, 32, 65536)] : index + %cols_b = index.assume %cols [range(%cols, 1, 16384)] : index + %total = index.mul %k_b, %cols_b : index + %groups0 = index.div %total, %c32 : index + %groups = index.assume %groups0 [range(%groups0, 1, 33554432)] : index + %groups64_0 = index.div %total, %c64 : index + %groups64 = index.assume %groups64_0 [range(%groups64_0, 1, 16777216)] : index + %nchunks = index.div %k_b, %c64 : index + + %wgid0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %wgid = index.assume %wgid0 [range(%wgid0, 0, 131071)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %wave = index.div %tid, %c32 : index + %lane = index.rem %tid, %c32 : index + %group_in_wave = index.div %lane, %c16 : index + %lane_in_group = index.rem %lane, %c16 : index + %group_half = index.div %lane_in_group, %c8 : index + %lane_in_half = index.rem %lane_in_group, %c8 : index + %is_scale_leader = index.cmp eq, %lane_in_group, %c0 : index + %is_sum_leader = index.cmp eq, %lane_in_half, %c0 : index + %wg_group_base = index.mul %wgid, %c128 : index + %wave_group_base = index.mul %wave, %c2 : index + %local_group_base = index.add %wg_group_base, %wave_group_base : index + + %src_g = buffer.assume.memory_space %src : buffer + %qs_g = buffer.assume.memory_space %qs : buffer + %ds_g = buffer.assume.memory_space %ds : buffer + %sums_g = buffer.assume.memory_space %sums : buffer + %src_na, %qs_na, %ds_na, %sums_na = buffer.assume.noalias %src_g, %qs_g, %ds_g, %sums_g : buffer, buffer, buffer, buffer + %src_elems = index.assume %total [range(%total, 32, 1073741824)] : index + %src_view = buffer.view %src_na[%base] : buffer -> view<[%src_elems]xf32> + %qs_halfwords = index.div %total, %c4 : index + %qs_view = buffer.view %qs_na[%base] : buffer -> view<[%qs_halfwords]xi16> + %ds_view = buffer.view %ds_na[%base] : buffer -> view<[%groups]xf32> + %sum_view = buffer.view %sums_na[%base] : buffer -> view<[%groups]xi32> + + scf.for %batch = [%c0 to %c8 step %c1] { + %batch_group_off = index.mul %batch, %c16 : index + %gid0 = index.add %local_group_base, %batch_group_off : index + %gid1 = index.add %gid0, %group_in_wave : index + %gid = index.assume %gid1 [range(%gid1, 0, 16777215)] : index + %in_range = index.cmp ult, %gid, %groups64 : index + scf.if %in_range { + %ebase0 = index.mul %gid, %c64 : index + %lane_elem_off = index.mul %lane_in_group, %c4 : index + %ebase1 = index.add %ebase0, %lane_elem_off : index + %ebase = index.assume %ebase1 [range(%ebase1, 0, 1073741820)] : index + %v = vector.load %src_view[%ebase] : view<[%src_elems]xf32> -> vector<4xf32> + %av = vector.absf %v : vector<4xf32> + %lane_max = vector.reduce %av, %c0_f32 : vector<4xf32>, f32 + %max_x1_peer, %max_x1_valid = kernel.subgroup.shuffle %lane_max, %xor1, %shuffle_width : f32, i32, i32 + %max_x1 = scalar.maxnumf %lane_max, %max_x1_peer : f32 + %max_x2_peer, %max_x2_valid = kernel.subgroup.shuffle %max_x1, %xor2, %shuffle_width : f32, i32, i32 + %max_x2 = scalar.maxnumf %max_x1, %max_x2_peer : f32 + %max_x4_peer, %max_x4_valid = kernel.subgroup.shuffle %max_x2, %xor4, %shuffle_width : f32, i32, i32 + %max_x4 = scalar.maxnumf %max_x2, %max_x4_peer : f32 + %max_x8_peer, %max_x8_valid = kernel.subgroup.shuffle %max_x4, %xor8, %shuffle_width : f32, i32, i32 + %group_max = scalar.maxnumf %max_x4, %max_x8_peer : f32 + %amax = scalar.maxnumf %group_max, %amax_eps : f32 + %a_scale = scalar.mulf %amax, %inv7 : f32 + %a_rscale = scalar.divf %f7, %amax : f32 + %gid32_base = index.mul %gid, %c2 : index + %gid32_high = index.add %gid32_base, %c1 : index + %token = index.div %gid, %nchunks : index + %chunk = index.rem %gid, %nchunks : index + %plane_chunk32 = index.mul %chunk, %c2 : index + %plane_low_base = index.mul %plane_chunk32, %cols_b : index + %plane_gid32_low = index.add %plane_low_base, %token : index + %plane_chunk32_high = index.add %plane_chunk32, %c1 : index + %plane_high_base = index.mul %plane_chunk32_high, %cols_b : index + %plane_gid32_high = index.add %plane_high_base, %token : index + %scale_low_index = scf.if %plane_major -> (index) { + scf.yield %plane_gid32_low : index + } else { + scf.yield %gid32_base : index + } + %scale_high_index = scf.if %plane_major -> (index) { + scf.yield %plane_gid32_high : index + } else { + scf.yield %gid32_high : index + } + scf.if %is_scale_leader { + view.store %a_scale, %ds_view[%scale_low_index] : f32, view<[%groups]xf32> + view.store %a_scale, %ds_view[%scale_high_index] : f32, view<[%groups]xf32> + } + + %rs = vector.splat %a_rscale : vector<4xf32> + %scaled = vector.mulf %v, %rs : vector<4xf32> + %rounded = vector.roundf %scaled : vector<4xf32> + %q32 = vector.fptosi %rounded : vector<4xf32> to vector<4xi32> + %qn = vector.andi %q32, %mask15 : vector<4xi32> + %q0 = vector.extract %qn[0] : vector<4xi32> -> i32 + %q1 = vector.extract %qn[1] : vector<4xi32> -> i32 + %q2 = vector.extract %qn[2] : vector<4xi32> -> i32 + %q3 = vector.extract %qn[3] : vector<4xi32> -> i32 + %q1s = scalar.shli %q1, %shift4 : i32 + %q2s = scalar.shli %q2, %shift8 : i32 + %q3s = scalar.shli %q3, %shift12 : i32 + %q01 = scalar.ori %q0, %q1s : i32 + %q23 = scalar.ori %q2s, %q3s : i32 + %qpacked32 = scalar.ori %q01, %q23 : i32 + %qpacked = scalar.trunci %qpacked32 : i32 to i16 + %plane_chunk_base = index.mul %chunk, %cols_b : index + %plane_gid = index.add %plane_chunk_base, %token : index + %payload_gid = scf.if %plane_major -> (index) { + scf.yield %plane_gid : index + } else { + scf.yield %gid : index + } + %word_base = index.mul %payload_gid, %c16 : index + %word_index0 = index.add %word_base, %lane_in_group : index + %word_index = index.assume %word_index0 [range(%word_index0, 0, 268435455)] : index + view.store %qpacked, %qs_view[%word_index] : i16, view<[%qs_halfwords]xi16> + + %lane_sum = vector.reduce %rounded, %c0_f32 : vector<4xf32>, f32 + %sum_x1_peer, %sum_x1_valid = kernel.subgroup.shuffle %lane_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum_x1 = scalar.addf %lane_sum, %sum_x1_peer : f32 + %sum_x2_peer, %sum_x2_valid = kernel.subgroup.shuffle %sum_x1, %xor2, %shuffle_width : f32, i32, i32 + %sum_x2 = scalar.addf %sum_x1, %sum_x2_peer : f32 + %sum_x4_peer, %sum_x4_valid = kernel.subgroup.shuffle %sum_x2, %xor4, %shuffle_width : f32, i32, i32 + %group_sum = scalar.addf %sum_x2, %sum_x4_peer : f32 + scf.if %is_sum_leader { + %sum_value = scalar.fptosi %group_sum : f32 to i32 + %row_sum_index = index.add %gid32_base, %group_half : index + %plane_sum_chunk = index.add %plane_chunk32, %group_half : index + %plane_sum_base = index.mul %plane_sum_chunk, %cols_b : index + %plane_sum_index = index.add %plane_sum_base, %token : index + %sum_index = scf.if %plane_major -> (index) { + scf.yield %plane_sum_index : index + } else { + scf.yield %row_sum_index : index + } + view.store %sum_value, %sum_view[%sum_index] : i32, view<[%groups]xi32> + } + } + } + template.return +} + +// --- gate/up contraction --- +template.decl @ggml.mul_mat_swiglu.symmetric_i4.m128n64_wg128.body(%apply_swiglu: i1, %publish_f32: i1, %publish_f16: i1, %publish_u4: i1, %publish_q8_plane: i1, %gate_na: buffer, %src0_na: buffer, %dst_na: buffer, %aq_na: buffer, %as_na: buffer, %asum_na: buffer, %qout_qs_na: buffer, %qout_ds_na: buffer, %qout_sums_na: buffer, %qout_scratch: buffer, %base: offset, %qact_scale_base: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index) + +kernel.def export("ggml_mul_mat_swiglu_symmetric_i4_wmma_q8_plane") @ggml_mul_mat_swiglu_symmetric_i4_wmma_q8_plane() { + %unit = index.constant 1 : index + %k_n = index.constant 64 : index + %k_m = index.constant 128 : index + %n_m1 = index.constant 63 : index + %m_m1 = index.constant 127 : index + %wg = index.constant 128 : index + %rows = config.get @ggml.mul_mat_swiglu.symmetric_i4.output_size : index + %cols = config.get @ggml.mul_mat_swiglu.symmetric_i4.token_count : index + %rows_up = index.add %rows, %n_m1 : index + %row_tiles = index.div %rows_up, %k_n : index + %cols_up = index.add %cols, %m_m1 : index + %col_tiles = index.div %cols_up, %k_m : index + kernel.launch.config workgroups(%col_tiles, %row_tiles, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%gate_weight: buffer, %up_weight: buffer, %input: buffer, %dst: buffer, %aq: buffer, %as: buffer, %asum: buffer, %q8_output: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 40 : index + %k_wstride = index.constant 40 : index + %k_m = index.constant 128 : index + %k_n = index.constant 64 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 4 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %apply_swiglu = scalar.constant true : i1 + %publish_f32 = scalar.constant false : i1 + %publish_f16 = scalar.constant false : i1 + %publish_u4 = scalar.constant false : i1 + %publish_q8_plane = scalar.constant true : i1 + %scratch_bytes = index.constant 16384 : offset + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + + %k = config.get @ggml.mul_mat_swiglu.symmetric_i4.input_size : index + %rows = config.get @ggml.mul_mat_swiglu.symmetric_i4.output_size : index + %cols = config.get @ggml.mul_mat_swiglu.symmetric_i4.token_count : index + %k_b = index.assume %k [range(%k, 64, 32768)] : index + %rows_b = index.assume %rows [range(%rows, 64, 262144), mul(%rows, 64)] : index + %cols_b = index.assume %cols [range(%cols, 128, 32768), mul(%cols, 128)] : index + %nchunks = index.div %k_b, %k_kc : index + %kblocks = index.div %k_b, %c32 : index + %col_tile0 = kernel.workgroup.id : index + %row_tile0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 255)] : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 8191)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 127)] : index + %wave = index.div %tid, %k_wave : index + + %gate_g = buffer.assume.memory_space %gate_weight : buffer + %up_g = buffer.assume.memory_space %up_weight : buffer + %dst_g = buffer.assume.memory_space %dst : buffer + %aq_g = buffer.assume.memory_space %aq : buffer + %as_g = buffer.assume.memory_space %as : buffer + %asum_g = buffer.assume.memory_space %asum : buffer + %q8_output_g = buffer.assume.memory_space %q8_output : buffer + %gate_na, %up_na, %dst_na, %aq_na, %as_na, %asum_na, %q8_output_na = buffer.assume.noalias %gate_g, %up_g, %dst_g, %aq_g, %as_g, %asum_g, %q8_output_g : buffer, buffer, buffer, buffer, buffer, buffer, buffer + + template.apply<@ggml.mul_mat_swiglu.symmetric_i4.m128n64_wg128.body>(%apply_swiglu, %publish_f32, %publish_f16, %publish_u4, %publish_q8_plane, %up_na, %gate_na, %dst_na, %aq_na, %as_na, %asum_na, %q8_output_na, %dst_na, %dst_na, %scratch, %base, %base, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %k_b, %rows_b, %cols_b, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave) : (i1, i1, i1, i1, i1, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index) + kernel.return +} + +config.decl @ggml.mul_mat_swiglu.symmetric_i4.input_size : %value: index where [range(%value, 32, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat_swiglu.symmetric_i4.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat_swiglu.symmetric_i4.token_count : %value: index where [range(%value, 1, 32768)] + +template.def<@ggml.mul_mat_swiglu.symmetric_i4.m128n64_wg128.body> device @ggml_mul_mat_swiglu_symmetric_i4_m128n64_wg128_body(%apply_swiglu: i1, %publish_f32: i1, %publish_f16: i1, %publish_u4: i1, %publish_q8_plane: i1, %gate_na: buffer, %src0_na: buffer, %dst_na: buffer, %aq_na: buffer, %as_na: buffer, %asum_na: buffer, %qout_qs_na: buffer, %qout_ds_na: buffer, %qout_sums_na: buffer, %qout_scratch: buffer, %base: offset, %qact_scale_base: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index) { + %aq_flat = buffer.view %aq_na[%base] : buffer -> view<1073741824xi8> + %as_flat = buffer.view %as_na[%qact_scale_base] : buffer -> view<33554432xf32> + %q4_k_block = index.constant 256 : index + // dst is [rows, cols] with rows contiguous, and C is [cols, rows]. + %dst_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + %dst_f16_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf16> + + // Packed I4 operands use 32 payload bytes per logical K64 row. + // A 40-byte LDS stride retains the conflict-avoiding padding of the IU8 body. + %scratch = buffer.assume.memory_space %qout_scratch : buffer + %i4_schema = encoding.define #encoding.operand : encoding + %al_flat = buffer.view %scratch[%base] : buffer -> view<5120xi8> + %w_off = index.constant 5120 : offset + %wl_flat = buffer.view %scratch[%w_off] : buffer -> view<5120xi8> + %as_off = index.constant 10240 : offset + %asl_view = buffer.view %scratch[%as_off] : buffer -> view<256xf32> + %ws_off = index.constant 11264 : offset + %wsl_view = buffer.view %scratch[%ws_off] : buffer -> view<256xf32> + %col_base = index.mul %col_tile, %k_m : index + %row_base = index.mul %row_tile, %k_n : index + %k_asn = index.constant 256 : index + %k_lanes0 = index.constant 0 : index + %k_lanes1 = index.constant 128 : index + + // Four waves own M32xN64 from gate and M32xN64 from up. + %wcol = index.rem %wave, %k_wm : index + %wrow = index.div %wave, %k_wm : index + %k_mspan = index.constant 16 : index + %k_nspan = index.constant 64 : index + %wm_off = index.mul %wcol, %k_mspan : index + %wn_off = index.mul %wrow, %k_nspan : index + %m_out = index.add %col_base, %wm_off : index + %n_out = index.add %row_base, %wn_off : index + // RDNA3 WMMA RHS lanes map to column lane % 16, so each lane starts at its output row. + %lane = index.rem %tid, %k_wave : index + %lane_lo = index.rem %lane, %c16 : index + %lane_hi = index.div %lane, %c16 : index + %wn_lane = index.add %wn_off, %lane_lo : index + %lm0 = index.add %wm_off, %c0 : index + %gm0 = index.add %m_out, %c0 : index + %k_ma1 = index.constant 16 : index + %lm1 = index.add %wm_off, %c0 : index + %gm1 = index.add %m_out, %c0 : index + %k_ma2 = index.constant 64 : index + %lm2 = index.add %wm_off, %k_ma2 : index + %gm2 = index.add %m_out, %k_ma2 : index + %lm3 = index.add %wm_off, %k_ma2 : index + %gm3 = index.add %m_out, %k_ma2 : index + %ln0 = index.add %wn_lane, %c0 : index + %gn0 = index.add %n_out, %c0 : index + %k_nb1 = index.constant 16 : index + %ln1 = index.add %wn_lane, %k_nb1 : index + %gn1 = index.add %n_out, %k_nb1 : index + %k_nb2 = index.constant 32 : index + %ln2 = index.add %wn_lane, %k_nb2 : index + %gn2 = index.add %n_out, %k_nb2 : index + %k_nb3 = index.constant 48 : index + %ln3 = index.add %wn_lane, %k_nb3 : index + %gn3 = index.add %n_out, %k_nb3 : index + %k_nb4 = index.constant 64 : index + %ln4 = index.add %wn_lane, %k_nb4 : index + %k_nb5 = index.constant 80 : index + %ln5 = index.add %wn_lane, %k_nb5 : index + %k_nb6 = index.constant 96 : index + %ln6 = index.add %wn_lane, %k_nb6 : index + %k_nb7 = index.constant 112 : index + %ln7 = index.add %wn_lane, %k_nb7 : index + %wsb0_0 = index.add %ln0, %c0 : index + %wsb0 = index.mul %wsb0_0, %k_blocks : index + %wsb1_0 = index.add %ln1, %c0 : index + %wsb1 = index.mul %wsb1_0, %k_blocks : index + %wsb2_0 = index.add %ln2, %c0 : index + %wsb2 = index.mul %wsb2_0, %k_blocks : index + %wsb3_0 = index.add %ln3, %c0 : index + %wsb3 = index.mul %wsb3_0, %k_blocks : index + %wsb4_0 = index.add %ln4, %c0 : index + %wsb4 = index.mul %wsb4_0, %k_blocks : index + %wsb5_0 = index.add %ln5, %c0 : index + %wsb5 = index.mul %wsb5_0, %k_blocks : index + %wsb6_0 = index.add %ln6, %c0 : index + %wsb6 = index.mul %wsb6_0, %k_blocks : index + %wsb7_0 = index.add %ln7, %c0 : index + %wsb7 = index.mul %wsb7_0, %k_blocks : index + // Activation scale metadata is [tile][block][lane-half][register]. + %as_tile0 = index.div %lm0, %c16 : index + %asb0 = index.mul %as_tile0, %k_blocks : index + %as_tile1 = index.div %lm1, %c16 : index + %asb1 = index.mul %as_tile1, %k_blocks : index + %as_tile2 = index.div %lm2, %c16 : index + %asb2 = index.mul %as_tile2, %k_blocks : index + %as_tile3 = index.div %lm3, %c16 : index + %asb3 = index.mul %as_tile3, %k_blocks : index + + %f0_0, %f0_1, %f0_2, %f0_3, %f1_0, %f1_1, %f1_2, %f1_3, %f2_0, %f2_1, %f2_2, %f2_3, %f3_0, %f3_1, %f3_2, %f3_3 = scf.for %chunk = [%c0 to %nchunks step %c1](%fc0_0 = %fzero : vector<8xf32>, %fc0_1 = %fzero : vector<8xf32>, %fc0_2 = %fzero : vector<8xf32>, %fc0_3 = %fzero : vector<8xf32>, %fc1_0 = %fzero : vector<8xf32>, %fc1_1 = %fzero : vector<8xf32>, %fc1_2 = %fzero : vector<8xf32>, %fc1_3 = %fzero : vector<8xf32>, %fc2_0 = %fzero : vector<8xf32>, %fc2_1 = %fzero : vector<8xf32>, %fc2_2 = %fzero : vector<8xf32>, %fc2_3 = %fzero : vector<8xf32>, %fc3_0 = %fzero : vector<8xf32>, %fc3_1 = %fzero : vector<8xf32>, %fc3_2 = %fzero : vector<8xf32>, %fc3_3 = %fzero : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %chunk_k = index.mul %chunk, %k_kc : index + %chunk_blk = index.mul %chunk, %k_blocks : index + + // All lanes stage the unique M128 activation rows. + %q4a_row0 = index.add %col_base, %tid : index + %q4a_row = index.assume %q4a_row0 [range(%q4a_row0, 0, 32767)] : index + %q4a_row_bytes = index.div %k_b, %c2 : index + %q4a_row_base = index.mul %q4a_row, %q4a_row_bytes : index + %q4a_chunk_bytes = index.div %chunk_k, %c2 : index + %q4a_src0_0 = index.add %q4a_row_base, %q4a_chunk_bytes : index + %q4a_src0 = index.assume %q4a_src0_0 [range(%q4a_src0_0, 0, 536870896)] : index + %q4a_v0 = vector.load %aq_flat[%q4a_src0] : view<1073741824xi8> -> vector<16xi8> + %q4a_src1_0 = index.add %q4a_src0, %c16 : index + %q4a_src1 = index.assume %q4a_src1_0 [range(%q4a_src1_0, 16, 536870912)] : index + %q4a_v1 = vector.load %aq_flat[%q4a_src1] : view<1073741824xi8> -> vector<16xi8> + %q4a_dst0_0 = index.mul %tid, %k_astride : index + %q4a_dst0 = index.assume %q4a_dst0_0 [range(%q4a_dst0_0, 0, 5080)] : index + vector.store %q4a_v0, %al_flat[%q4a_dst0] : vector<16xi8>, view<5120xi8> + %q4a_dst1_0 = index.add %q4a_dst0, %c16 : index + %q4a_dst1 = index.assume %q4a_dst1_0 [range(%q4a_dst1_0, 16, 5096)] : index + vector.store %q4a_v1, %al_flat[%q4a_dst1] : vector<16xi8>, view<5120xi8> + + // The 144-byte block holds eight duplicated K64 scale slots followed by eight signed-I4 K32 payloads. + %q4_weight_lane_count = index.constant 128 : index + %q4_weight_lane_active = index.cmp ult, %tid, %q4_weight_lane_count : index + scf.if %q4_weight_lane_active { + %q4_projection_rows = index.constant 64 : index + %q4_weight_tid0 = index.rem %tid, %q4_projection_rows : index + %q4_weight_tid = index.assume %q4_weight_tid0 [range(%q4_weight_tid0, 0, 63)] : index + %q4_lds_row = index.assume %tid [range(%tid, 0, 127)] : index + %q4_is_up = index.cmp uge, %tid, %q4_projection_rows : index + %q4_block_count = index.div %k_b, %q4_k_block : index + %q4_block = index.div %chunk, %c4 : index + %q4_pair = index.rem %chunk, %c4 : index + %q4_row0 = index.add %row_base, %q4_weight_tid : index + %q4_row = index.assume %q4_row0 [range(%q4_row0, 0, 262143)] : index + %q4_row_group_size = index.constant 64 : index + %q4_field_count = index.constant 9 : index + %q4_field_bytes = index.constant 16 : index + %q4_row_group = index.div %q4_row, %q4_row_group_size : index + %q4_row_lane = index.rem %q4_row, %q4_row_group_size : index + %q4_group_block0 = index.mul %q4_row_group, %q4_block_count : index + %q4_group_block = index.add %q4_group_block0, %q4_block : index + %q4_group_field_base = index.mul %q4_group_block, %q4_field_count : index + %q4_header_lane_base = index.mul %q4_group_field_base, %q4_row_group_size : index + %q4_header_lane = index.add %q4_header_lane_base, %q4_row_lane : index + %q4_header_byte_index = index.mul %q4_header_lane, %q4_field_bytes : index + %q4_scale_byte_stride = index.constant 4 : index + %q4_scale_byte_offset = index.mul %q4_pair, %q4_scale_byte_stride : index + %q4_scale_byte_index = index.add %q4_header_byte_index, %q4_scale_byte_offset : index + %q4_scale_byte_base = index.cast %q4_scale_byte_index : index to offset + %q4_scale_pair = scf.if %q4_is_up -> (vector<2xf16>) { + %q4_up_scale_view = buffer.view %gate_na[%q4_scale_byte_base] : buffer -> view<2xf16> + %q4_up_scale_pair = vector.load %q4_up_scale_view[%c0] : view<2xf16> -> vector<2xf16> + scf.yield %q4_up_scale_pair : vector<2xf16> + } else { + %q4_gate_scale_view = buffer.view %src0_na[%q4_scale_byte_base] : buffer -> view<2xf16> + %q4_gate_scale_pair = vector.load %q4_gate_scale_view[%c0] : view<2xf16> -> vector<2xf16> + scf.yield %q4_gate_scale_pair : vector<2xf16> + } + %q4_scale_low_f16 = vector.extract %q4_scale_pair[0] : vector<2xf16> -> f16 + %q4_scale_high_f16 = vector.extract %q4_scale_pair[1] : vector<2xf16> -> f16 + %q4_d0 = scalar.extf %q4_scale_low_f16 : f16 to f32 + %q4_d1 = scalar.extf %q4_scale_high_f16 : f16 to f32 + %q4_low_group = index.mul %q4_pair, %c2 : index + %q4_high_group = index.add %q4_low_group, %c1 : index + %q4_low_field0 = index.add %q4_low_group, %c1 : index + %q4_high_field0 = index.add %q4_high_group, %c1 : index + %q4_low_field = index.add %q4_group_field_base, %q4_low_field0 : index + %q4_high_field = index.add %q4_group_field_base, %q4_high_field0 : index + %q4_low_lane_base = index.mul %q4_low_field, %q4_row_group_size : index + %q4_high_lane_base = index.mul %q4_high_field, %q4_row_group_size : index + %q4_low_lane = index.add %q4_low_lane_base, %q4_row_lane : index + %q4_high_lane = index.add %q4_high_lane_base, %q4_row_lane : index + %q4_low_byte_index = index.mul %q4_low_lane, %q4_field_bytes : index + %q4_high_byte_index = index.mul %q4_high_lane, %q4_field_bytes : index + %q4_low_byte_base = index.cast %q4_low_byte_index : index to offset + %q4_high_byte_base = index.cast %q4_high_byte_index : index to offset + %q4_low_words, %q4_high_words = scf.if %q4_is_up -> (vector<4xi32>, vector<4xi32>) { + %q4_up_low_view = buffer.view %gate_na[%q4_low_byte_base] : buffer -> view<4xi32> + %q4_up_high_view = buffer.view %gate_na[%q4_high_byte_base] : buffer -> view<4xi32> + %q4_up_low_words = vector.load %q4_up_low_view[%c0] : view<4xi32> -> vector<4xi32> + %q4_up_high_words = vector.load %q4_up_high_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %q4_up_low_words, %q4_up_high_words : vector<4xi32>, vector<4xi32> + } else { + %q4_gate_low_view = buffer.view %src0_na[%q4_low_byte_base] : buffer -> view<4xi32> + %q4_gate_high_view = buffer.view %src0_na[%q4_high_byte_base] : buffer -> view<4xi32> + %q4_gate_low_words = vector.load %q4_gate_low_view[%c0] : view<4xi32> -> vector<4xi32> + %q4_gate_high_words = vector.load %q4_gate_high_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %q4_gate_low_words, %q4_gate_high_words : vector<4xi32>, vector<4xi32> + } + %q4_low_packed = vector.bitcast %q4_low_words : vector<4xi32> to vector<16xi8> + %q4_high_packed = vector.bitcast %q4_high_words : vector<4xi32> to vector<16xi8> + %q4_weight_row_base = index.mul %q4_lds_row, %k_wstride : index + %q4_w0 = index.assume %q4_weight_row_base [range(%q4_weight_row_base, 0, 5080)] : index + %q4_w1_0 = index.add %q4_weight_row_base, %c16 : index + %q4_w1 = index.assume %q4_w1_0 [range(%q4_w1_0, 16, 5096)] : index + vector.store %q4_low_packed, %wl_flat[%q4_w0] : vector<16xi8>, view<5120xi8> + vector.store %q4_high_packed, %wl_flat[%q4_w1] : vector<16xi8>, view<5120xi8> + %q4_meta0_0 = index.mul %q4_lds_row, %c2 : index + %q4_meta0 = index.assume %q4_meta0_0 [range(%q4_meta0_0, 0, 254)] : index + %q4_meta1_0 = index.add %q4_meta0, %c1 : index + %q4_meta1 = index.assume %q4_meta1_0 [range(%q4_meta1_0, 1, 255)] : index + view.store %q4_d0, %wsl_view[%q4_meta0] : f32, view<256xf32> + view.store %q4_d1, %wsl_view[%q4_meta1] : f32, view<256xf32> + } + // Scale planes: 256 activation and 128 weight entries. + %sca0_slot0 = index.add %tid, %k_lanes0 : index + %sca0_slot = index.assume %sca0_slot0 [range(%sca0_slot0, 0, 127)] : index + %scin_a0 = index.cmp ult, %sca0_slot, %k_asn : index + scf.if %scin_a0 { + %sc0_col = index.div %sca0_slot, %k_blocks : index + %sc0_blk = index.rem %sca0_slot, %k_blocks : index + %sca0_c = index.add %col_base, %sc0_col : index + %sca0_0 = index.mul %sca0_c, %kblocks : index + %sca0_1 = index.add %sca0_0, %chunk_blk : index + %sca0_2 = index.add %sca0_1, %sc0_blk : index + %sca0 = index.assume %sca0_2 [range(%sca0_2, 0, 33554431)] : index + %scav0 = view.load %as_flat[%sca0] : view<33554432xf32> -> f32 + %sca0_tile = index.div %sc0_col, %c16 : index + %sca0_inner = index.rem %sc0_col, %c16 : index + %sca0_par = index.rem %sca0_inner, %c2 : index + %sca0_v = index.div %sca0_inner, %c2 : index + %sca0_dst0 = index.mul %sca0_tile, %k_blocks : index + %sca0_dst1 = index.add %sca0_dst0, %sc0_blk : index + %sca0_dst2 = index.mul %sca0_dst1, %c2 : index + %sca0_dst3 = index.add %sca0_dst2, %sca0_par : index + %sca0_dst4 = index.mul %sca0_dst3, %c8 : index + %sca0_dst5 = index.add %sca0_dst4, %sca0_v : index + %sca0_dst = index.assume %sca0_dst5 [range(%sca0_dst5, 0, 255)] : index + view.store %scav0, %asl_view[%sca0_dst] : f32, view<256xf32> + } + %sca1_slot0 = index.add %tid, %k_lanes1 : index + %sca1_slot = index.assume %sca1_slot0 [range(%sca1_slot0, 128, 255)] : index + %scin_a1 = index.cmp ult, %sca1_slot, %k_asn : index + scf.if %scin_a1 { + %sc1_col = index.div %sca1_slot, %k_blocks : index + %sc1_blk = index.rem %sca1_slot, %k_blocks : index + %sca1_c = index.add %col_base, %sc1_col : index + %sca1_0 = index.mul %sca1_c, %kblocks : index + %sca1_1 = index.add %sca1_0, %chunk_blk : index + %sca1_2 = index.add %sca1_1, %sc1_blk : index + %sca1 = index.assume %sca1_2 [range(%sca1_2, 0, 33554431)] : index + %scav1 = view.load %as_flat[%sca1] : view<33554432xf32> -> f32 + %sca1_tile = index.div %sc1_col, %c16 : index + %sca1_inner = index.rem %sc1_col, %c16 : index + %sca1_par = index.rem %sca1_inner, %c2 : index + %sca1_v = index.div %sca1_inner, %c2 : index + %sca1_dst0 = index.mul %sca1_tile, %k_blocks : index + %sca1_dst1 = index.add %sca1_dst0, %sc1_blk : index + %sca1_dst2 = index.mul %sca1_dst1, %c2 : index + %sca1_dst3 = index.add %sca1_dst2, %sca1_par : index + %sca1_dst4 = index.mul %sca1_dst3, %c8 : index + %sca1_dst5 = index.add %sca1_dst4, %sca1_v : index + %sca1_dst = index.assume %sca1_dst5 [range(%sca1_dst5, 0, 255)] : index + view.store %scav1, %asl_view[%sca1_dst] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + // Two K32 blocks explicitly unrolled. Arithmetic order and the + // enclosing K64 staging/barrier cadence are unchanged. + %blk_k_u0 = index.mul %c0, %c16 : index + %blk_k16_u0 = index.add %blk_k_u0, %c8 : index + %iz_u0 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %wsi0_2_u0 = index.add %wsb0, %c0 : index + %wsi0_u0 = index.assume %wsi0_2_u0 [range(%wsi0_2_u0, 0, 255)] : index + %wsc0_u0 = view.load %wsl_view[%wsi0_u0] : view<256xf32> -> f32 + %wsv0_u0 = vector.splat %wsc0_u0 : vector<8xf32> + %wsi1_2_u0 = index.add %wsb1, %c0 : index + %wsi1_u0 = index.assume %wsi1_2_u0 [range(%wsi1_2_u0, 0, 255)] : index + %wsc1_u0 = view.load %wsl_view[%wsi1_u0] : view<256xf32> -> f32 + %wsv1_u0 = vector.splat %wsc1_u0 : vector<8xf32> + %wsi4_2_u0 = index.add %wsb4, %c0 : index + %wsi4_u0 = index.assume %wsi4_2_u0 [range(%wsi4_2_u0, 0, 255)] : index + %wsc4_u0 = view.load %wsl_view[%wsi4_u0] : view<256xf32> -> f32 + %wsv4_u0 = vector.splat %wsc4_u0 : vector<8xf32> + %wsi5_2_u0 = index.add %wsb5, %c0 : index + %wsi5_u0 = index.assume %wsi5_2_u0 [range(%wsi5_2_u0, 0, 255)] : index + %wsc5_u0 = view.load %wsl_view[%wsi5_u0] : view<256xf32> -> f32 + %wsv5_u0 = vector.splat %wsc5_u0 : vector<8xf32> + %lf0_0_lane_row_u0 = index.add %lm0, %lane_lo : index + %lf0_0_row_u0 = index.mul %lf0_0_lane_row_u0, %k_astride : index + %lf0_0_base_u0 = index.add %lf0_0_row_u0, %blk_k_u0 : index + %lf0_0_idx_u0 = index.assume %lf0_0_base_u0 [range(%lf0_0_base_u0, 0, 5112)] : index + %lf0_0_raw_u0 = vector.load %al_flat[%lf0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf0_0_words_u0 = vector.bitcast %lf0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %lf0_0_u0 = vector.fragment %lf0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf0_1_lane_row_u0 = index.add %lm0, %lane_lo : index + %lf0_1_row_u0 = index.mul %lf0_1_lane_row_u0, %k_astride : index + %lf0_1_base_u0 = index.add %lf0_1_row_u0, %blk_k16_u0 : index + %lf0_1_idx_u0 = index.assume %lf0_1_base_u0 [range(%lf0_1_base_u0, 0, 5112)] : index + %lf0_1_raw_u0 = vector.load %al_flat[%lf0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf0_1_words_u0 = vector.bitcast %lf0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %lf0_1_u0 = vector.fragment %lf0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %asi0_0_u0 = index.add %asb0, %c0 : index + %asi0_1_u0 = index.mul %asi0_0_u0, %c2 : index + %asi0_2_u0 = index.add %asi0_1_u0, %lane_hi : index + %asi0_3_u0 = index.mul %asi0_2_u0, %c8 : index + %asi0_u0 = index.assume %asi0_3_u0 [range(%asi0_3_u0, 0, 120)] : index + %asv0_u0 = vector.load %asl_view[%asi0_u0] : view<256xf32> -> vector<8xf32> + %rfs0_0_0_row_u0 = index.mul %ln0, %k_wstride : index + %rfs0_0_0_base_u0 = index.add %rfs0_0_0_row_u0, %blk_k_u0 : index + %rfs0_0_0_idx_u0 = index.assume %rfs0_0_0_base_u0 [range(%rfs0_0_0_base_u0, 0, 5112)] : index + %rfs0_0_0_raw_u0 = vector.load %wl_flat[%rfs0_0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_0_0_words_u0 = vector.bitcast %rfs0_0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_0_0_u0 = vector.fragment %rfs0_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_1_row_u0 = index.mul %ln0, %k_wstride : index + %rfs0_0_1_base_u0 = index.add %rfs0_0_1_row_u0, %blk_k16_u0 : index + %rfs0_0_1_idx_u0 = index.assume %rfs0_0_1_base_u0 [range(%rfs0_0_1_base_u0, 0, 5112)] : index + %rfs0_0_1_raw_u0 = vector.load %wl_flat[%rfs0_0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_0_1_words_u0 = vector.bitcast %rfs0_0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_0_1_u0 = vector.fragment %rfs0_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_0_u0 = vector.mma %lf0_0_u0, %rfs0_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_0_u0 = vector.mma %lf0_1_u0, %rfs0_0_1_u0, %i0_0_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_0_u0 = vector.sitofp %i1_0_0_u0 : vector<8xi32> to vector<8xf32> + %sv0_0_u0 = vector.mulf %asv0_u0, %wsv0_u0 : vector<8xf32> + %fn0_0_u0 = vector.fmaf %ff0_0_u0, %sv0_0_u0, %fc0_0 : vector<8xf32> + scf.schedule.fence + + %rfs0_1_0_row_u0 = index.mul %ln1, %k_wstride : index + %rfs0_1_0_base_u0 = index.add %rfs0_1_0_row_u0, %blk_k_u0 : index + %rfs0_1_0_idx_u0 = index.assume %rfs0_1_0_base_u0 [range(%rfs0_1_0_base_u0, 0, 5112)] : index + %rfs0_1_0_raw_u0 = vector.load %wl_flat[%rfs0_1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_1_0_words_u0 = vector.bitcast %rfs0_1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_1_0_u0 = vector.fragment %rfs0_1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_1_1_row_u0 = index.mul %ln1, %k_wstride : index + %rfs0_1_1_base_u0 = index.add %rfs0_1_1_row_u0, %blk_k16_u0 : index + %rfs0_1_1_idx_u0 = index.assume %rfs0_1_1_base_u0 [range(%rfs0_1_1_base_u0, 0, 5112)] : index + %rfs0_1_1_raw_u0 = vector.load %wl_flat[%rfs0_1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_1_1_words_u0 = vector.bitcast %rfs0_1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_1_1_u0 = vector.fragment %rfs0_1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_1_u0 = vector.mma %lf0_0_u0, %rfs0_1_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_1_u0 = vector.mma %lf0_1_u0, %rfs0_1_1_u0, %i0_0_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_1_u0 = vector.sitofp %i1_0_1_u0 : vector<8xi32> to vector<8xf32> + %sv0_1_u0 = vector.mulf %asv0_u0, %wsv1_u0 : vector<8xf32> + %fn0_1_u0 = vector.fmaf %ff0_1_u0, %sv0_1_u0, %fc0_1 : vector<8xf32> + scf.schedule.fence + + %rfs0_2_0_row_u0 = index.mul %ln4, %k_wstride : index + %rfs0_2_0_base_u0 = index.add %rfs0_2_0_row_u0, %blk_k_u0 : index + %rfs0_2_0_idx_u0 = index.assume %rfs0_2_0_base_u0 [range(%rfs0_2_0_base_u0, 0, 5112)] : index + %rfs0_2_0_raw_u0 = vector.load %wl_flat[%rfs0_2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_2_0_words_u0 = vector.bitcast %rfs0_2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_2_0_u0 = vector.fragment %rfs0_2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_2_1_row_u0 = index.mul %ln4, %k_wstride : index + %rfs0_2_1_base_u0 = index.add %rfs0_2_1_row_u0, %blk_k16_u0 : index + %rfs0_2_1_idx_u0 = index.assume %rfs0_2_1_base_u0 [range(%rfs0_2_1_base_u0, 0, 5112)] : index + %rfs0_2_1_raw_u0 = vector.load %wl_flat[%rfs0_2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_2_1_words_u0 = vector.bitcast %rfs0_2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_2_1_u0 = vector.fragment %rfs0_2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_2_u0 = vector.mma %lf0_0_u0, %rfs0_2_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_2_u0 = vector.mma %lf0_1_u0, %rfs0_2_1_u0, %i0_0_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_2_u0 = vector.sitofp %i1_0_2_u0 : vector<8xi32> to vector<8xf32> + %sv0_2_u0 = vector.mulf %asv0_u0, %wsv4_u0 : vector<8xf32> + %fn0_2_u0 = vector.fmaf %ff0_2_u0, %sv0_2_u0, %fc0_2 : vector<8xf32> + scf.schedule.fence + + %rfs0_3_0_row_u0 = index.mul %ln5, %k_wstride : index + %rfs0_3_0_base_u0 = index.add %rfs0_3_0_row_u0, %blk_k_u0 : index + %rfs0_3_0_idx_u0 = index.assume %rfs0_3_0_base_u0 [range(%rfs0_3_0_base_u0, 0, 5112)] : index + %rfs0_3_0_raw_u0 = vector.load %wl_flat[%rfs0_3_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_3_0_words_u0 = vector.bitcast %rfs0_3_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_3_0_u0 = vector.fragment %rfs0_3_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_3_1_row_u0 = index.mul %ln5, %k_wstride : index + %rfs0_3_1_base_u0 = index.add %rfs0_3_1_row_u0, %blk_k16_u0 : index + %rfs0_3_1_idx_u0 = index.assume %rfs0_3_1_base_u0 [range(%rfs0_3_1_base_u0, 0, 5112)] : index + %rfs0_3_1_raw_u0 = vector.load %wl_flat[%rfs0_3_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_3_1_words_u0 = vector.bitcast %rfs0_3_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_3_1_u0 = vector.fragment %rfs0_3_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_3_u0 = vector.mma %lf0_0_u0, %rfs0_3_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_3_u0 = vector.mma %lf0_1_u0, %rfs0_3_1_u0, %i0_0_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_3_u0 = vector.sitofp %i1_0_3_u0 : vector<8xi32> to vector<8xf32> + %sv0_3_u0 = vector.mulf %asv0_u0, %wsv5_u0 : vector<8xf32> + %fn0_3_u0 = vector.fmaf %ff0_3_u0, %sv0_3_u0, %fc0_3 : vector<8xf32> + + scf.schedule.fence + %wsi2_2_u0 = index.add %wsb2, %c0 : index + %wsi2_u0 = index.assume %wsi2_2_u0 [range(%wsi2_2_u0, 0, 255)] : index + %wsc2_u0 = view.load %wsl_view[%wsi2_u0] : view<256xf32> -> f32 + %wsv2_u0 = vector.splat %wsc2_u0 : vector<8xf32> + %wsi3_2_u0 = index.add %wsb3, %c0 : index + %wsi3_u0 = index.assume %wsi3_2_u0 [range(%wsi3_2_u0, 0, 255)] : index + %wsc3_u0 = view.load %wsl_view[%wsi3_u0] : view<256xf32> -> f32 + %wsv3_u0 = vector.splat %wsc3_u0 : vector<8xf32> + %wsi6_2_u0 = index.add %wsb6, %c0 : index + %wsi6_u0 = index.assume %wsi6_2_u0 [range(%wsi6_2_u0, 0, 255)] : index + %wsc6_u0 = view.load %wsl_view[%wsi6_u0] : view<256xf32> -> f32 + %wsv6_u0 = vector.splat %wsc6_u0 : vector<8xf32> + %wsi7_2_u0 = index.add %wsb7, %c0 : index + %wsi7_u0 = index.assume %wsi7_2_u0 [range(%wsi7_2_u0, 0, 255)] : index + %wsc7_u0 = view.load %wsl_view[%wsi7_u0] : view<256xf32> -> f32 + %wsv7_u0 = vector.splat %wsc7_u0 : vector<8xf32> + %rfs1_0_0_row_u0 = index.mul %ln2, %k_wstride : index + %rfs1_0_0_base_u0 = index.add %rfs1_0_0_row_u0, %blk_k_u0 : index + %rfs1_0_0_idx_u0 = index.assume %rfs1_0_0_base_u0 [range(%rfs1_0_0_base_u0, 0, 5112)] : index + %rfs1_0_0_raw_u0 = vector.load %wl_flat[%rfs1_0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_0_0_words_u0 = vector.bitcast %rfs1_0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_0_0_u0 = vector.fragment %rfs1_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_0_1_row_u0 = index.mul %ln2, %k_wstride : index + %rfs1_0_1_base_u0 = index.add %rfs1_0_1_row_u0, %blk_k16_u0 : index + %rfs1_0_1_idx_u0 = index.assume %rfs1_0_1_base_u0 [range(%rfs1_0_1_base_u0, 0, 5112)] : index + %rfs1_0_1_raw_u0 = vector.load %wl_flat[%rfs1_0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_0_1_words_u0 = vector.bitcast %rfs1_0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_0_1_u0 = vector.fragment %rfs1_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_0_u0 = vector.mma %lf0_0_u0, %rfs1_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_0_u0 = vector.mma %lf0_1_u0, %rfs1_0_1_u0, %i0_1_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_0_u0 = vector.sitofp %i1_1_0_u0 : vector<8xi32> to vector<8xf32> + %sv1_0_u0 = vector.mulf %asv0_u0, %wsv2_u0 : vector<8xf32> + %fn1_0_u0 = vector.fmaf %ff1_0_u0, %sv1_0_u0, %fc1_0 : vector<8xf32> + scf.schedule.fence + + %rfs1_1_0_row_u0 = index.mul %ln3, %k_wstride : index + %rfs1_1_0_base_u0 = index.add %rfs1_1_0_row_u0, %blk_k_u0 : index + %rfs1_1_0_idx_u0 = index.assume %rfs1_1_0_base_u0 [range(%rfs1_1_0_base_u0, 0, 5112)] : index + %rfs1_1_0_raw_u0 = vector.load %wl_flat[%rfs1_1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_1_0_words_u0 = vector.bitcast %rfs1_1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_1_0_u0 = vector.fragment %rfs1_1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_1_1_row_u0 = index.mul %ln3, %k_wstride : index + %rfs1_1_1_base_u0 = index.add %rfs1_1_1_row_u0, %blk_k16_u0 : index + %rfs1_1_1_idx_u0 = index.assume %rfs1_1_1_base_u0 [range(%rfs1_1_1_base_u0, 0, 5112)] : index + %rfs1_1_1_raw_u0 = vector.load %wl_flat[%rfs1_1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_1_1_words_u0 = vector.bitcast %rfs1_1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_1_1_u0 = vector.fragment %rfs1_1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_1_u0 = vector.mma %lf0_0_u0, %rfs1_1_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_1_u0 = vector.mma %lf0_1_u0, %rfs1_1_1_u0, %i0_1_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_1_u0 = vector.sitofp %i1_1_1_u0 : vector<8xi32> to vector<8xf32> + %sv1_1_u0 = vector.mulf %asv0_u0, %wsv3_u0 : vector<8xf32> + %fn1_1_u0 = vector.fmaf %ff1_1_u0, %sv1_1_u0, %fc1_1 : vector<8xf32> + scf.schedule.fence + + %rfs1_2_0_row_u0 = index.mul %ln6, %k_wstride : index + %rfs1_2_0_base_u0 = index.add %rfs1_2_0_row_u0, %blk_k_u0 : index + %rfs1_2_0_idx_u0 = index.assume %rfs1_2_0_base_u0 [range(%rfs1_2_0_base_u0, 0, 5112)] : index + %rfs1_2_0_raw_u0 = vector.load %wl_flat[%rfs1_2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_2_0_words_u0 = vector.bitcast %rfs1_2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_2_0_u0 = vector.fragment %rfs1_2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_2_1_row_u0 = index.mul %ln6, %k_wstride : index + %rfs1_2_1_base_u0 = index.add %rfs1_2_1_row_u0, %blk_k16_u0 : index + %rfs1_2_1_idx_u0 = index.assume %rfs1_2_1_base_u0 [range(%rfs1_2_1_base_u0, 0, 5112)] : index + %rfs1_2_1_raw_u0 = vector.load %wl_flat[%rfs1_2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_2_1_words_u0 = vector.bitcast %rfs1_2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_2_1_u0 = vector.fragment %rfs1_2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_2_u0 = vector.mma %lf0_0_u0, %rfs1_2_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_2_u0 = vector.mma %lf0_1_u0, %rfs1_2_1_u0, %i0_1_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_2_u0 = vector.sitofp %i1_1_2_u0 : vector<8xi32> to vector<8xf32> + %sv1_2_u0 = vector.mulf %asv0_u0, %wsv6_u0 : vector<8xf32> + %fn1_2_u0 = vector.fmaf %ff1_2_u0, %sv1_2_u0, %fc1_2 : vector<8xf32> + scf.schedule.fence + + %rfs1_3_0_row_u0 = index.mul %ln7, %k_wstride : index + %rfs1_3_0_base_u0 = index.add %rfs1_3_0_row_u0, %blk_k_u0 : index + %rfs1_3_0_idx_u0 = index.assume %rfs1_3_0_base_u0 [range(%rfs1_3_0_base_u0, 0, 5112)] : index + %rfs1_3_0_raw_u0 = vector.load %wl_flat[%rfs1_3_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_3_0_words_u0 = vector.bitcast %rfs1_3_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_3_0_u0 = vector.fragment %rfs1_3_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_3_1_row_u0 = index.mul %ln7, %k_wstride : index + %rfs1_3_1_base_u0 = index.add %rfs1_3_1_row_u0, %blk_k16_u0 : index + %rfs1_3_1_idx_u0 = index.assume %rfs1_3_1_base_u0 [range(%rfs1_3_1_base_u0, 0, 5112)] : index + %rfs1_3_1_raw_u0 = vector.load %wl_flat[%rfs1_3_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_3_1_words_u0 = vector.bitcast %rfs1_3_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_3_1_u0 = vector.fragment %rfs1_3_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_3_u0 = vector.mma %lf0_0_u0, %rfs1_3_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_3_u0 = vector.mma %lf0_1_u0, %rfs1_3_1_u0, %i0_1_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_3_u0 = vector.sitofp %i1_1_3_u0 : vector<8xi32> to vector<8xf32> + %sv1_3_u0 = vector.mulf %asv0_u0, %wsv7_u0 : vector<8xf32> + %fn1_3_u0 = vector.fmaf %ff1_3_u0, %sv1_3_u0, %fc1_3 : vector<8xf32> + + scf.schedule.fence + %lf2_0_lane_row_u0 = index.add %lm2, %lane_lo : index + %lf2_0_row_u0 = index.mul %lf2_0_lane_row_u0, %k_astride : index + %lf2_0_base_u0 = index.add %lf2_0_row_u0, %blk_k_u0 : index + %lf2_0_idx_u0 = index.assume %lf2_0_base_u0 [range(%lf2_0_base_u0, 0, 5112)] : index + %lf2_0_raw_u0 = vector.load %al_flat[%lf2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf2_0_words_u0 = vector.bitcast %lf2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %lf2_0_u0 = vector.fragment %lf2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf2_1_lane_row_u0 = index.add %lm2, %lane_lo : index + %lf2_1_row_u0 = index.mul %lf2_1_lane_row_u0, %k_astride : index + %lf2_1_base_u0 = index.add %lf2_1_row_u0, %blk_k16_u0 : index + %lf2_1_idx_u0 = index.assume %lf2_1_base_u0 [range(%lf2_1_base_u0, 0, 5112)] : index + %lf2_1_raw_u0 = vector.load %al_flat[%lf2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf2_1_words_u0 = vector.bitcast %lf2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %lf2_1_u0 = vector.fragment %lf2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %asi2_0_u0 = index.add %asb2, %c0 : index + %asi2_1_u0 = index.mul %asi2_0_u0, %c2 : index + %asi2_2_u0 = index.add %asi2_1_u0, %lane_hi : index + %asi2_3_u0 = index.mul %asi2_2_u0, %c8 : index + %asi2_u0 = index.assume %asi2_3_u0 [range(%asi2_3_u0, 128, 248)] : index + %asv2_u0 = vector.load %asl_view[%asi2_u0] : view<256xf32> -> vector<8xf32> + %rfs2_0_0_row_u0 = index.mul %ln0, %k_wstride : index + %rfs2_0_0_base_u0 = index.add %rfs2_0_0_row_u0, %blk_k_u0 : index + %rfs2_0_0_idx_u0 = index.assume %rfs2_0_0_base_u0 [range(%rfs2_0_0_base_u0, 0, 5112)] : index + %rfs2_0_0_raw_u0 = vector.load %wl_flat[%rfs2_0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_0_0_words_u0 = vector.bitcast %rfs2_0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_0_0_u0 = vector.fragment %rfs2_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_0_1_row_u0 = index.mul %ln0, %k_wstride : index + %rfs2_0_1_base_u0 = index.add %rfs2_0_1_row_u0, %blk_k16_u0 : index + %rfs2_0_1_idx_u0 = index.assume %rfs2_0_1_base_u0 [range(%rfs2_0_1_base_u0, 0, 5112)] : index + %rfs2_0_1_raw_u0 = vector.load %wl_flat[%rfs2_0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_0_1_words_u0 = vector.bitcast %rfs2_0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_0_1_u0 = vector.fragment %rfs2_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_0_u0 = vector.mma %lf2_0_u0, %rfs2_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_0_u0 = vector.mma %lf2_1_u0, %rfs2_0_1_u0, %i0_2_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff2_0_u0 = vector.sitofp %i1_2_0_u0 : vector<8xi32> to vector<8xf32> + %sv2_0_u0 = vector.mulf %asv2_u0, %wsv0_u0 : vector<8xf32> + %fn2_0_u0 = vector.fmaf %ff2_0_u0, %sv2_0_u0, %fc2_0 : vector<8xf32> + scf.schedule.fence + + %rfs2_1_0_row_u0 = index.mul %ln1, %k_wstride : index + %rfs2_1_0_base_u0 = index.add %rfs2_1_0_row_u0, %blk_k_u0 : index + %rfs2_1_0_idx_u0 = index.assume %rfs2_1_0_base_u0 [range(%rfs2_1_0_base_u0, 0, 5112)] : index + %rfs2_1_0_raw_u0 = vector.load %wl_flat[%rfs2_1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_1_0_words_u0 = vector.bitcast %rfs2_1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_1_0_u0 = vector.fragment %rfs2_1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_1_1_row_u0 = index.mul %ln1, %k_wstride : index + %rfs2_1_1_base_u0 = index.add %rfs2_1_1_row_u0, %blk_k16_u0 : index + %rfs2_1_1_idx_u0 = index.assume %rfs2_1_1_base_u0 [range(%rfs2_1_1_base_u0, 0, 5112)] : index + %rfs2_1_1_raw_u0 = vector.load %wl_flat[%rfs2_1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_1_1_words_u0 = vector.bitcast %rfs2_1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_1_1_u0 = vector.fragment %rfs2_1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_1_u0 = vector.mma %lf2_0_u0, %rfs2_1_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_1_u0 = vector.mma %lf2_1_u0, %rfs2_1_1_u0, %i0_2_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff2_1_u0 = vector.sitofp %i1_2_1_u0 : vector<8xi32> to vector<8xf32> + %sv2_1_u0 = vector.mulf %asv2_u0, %wsv1_u0 : vector<8xf32> + %fn2_1_u0 = vector.fmaf %ff2_1_u0, %sv2_1_u0, %fc2_1 : vector<8xf32> + scf.schedule.fence + + %rfs2_2_0_row_u0 = index.mul %ln4, %k_wstride : index + %rfs2_2_0_base_u0 = index.add %rfs2_2_0_row_u0, %blk_k_u0 : index + %rfs2_2_0_idx_u0 = index.assume %rfs2_2_0_base_u0 [range(%rfs2_2_0_base_u0, 0, 5112)] : index + %rfs2_2_0_raw_u0 = vector.load %wl_flat[%rfs2_2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_2_0_words_u0 = vector.bitcast %rfs2_2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_2_0_u0 = vector.fragment %rfs2_2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_2_1_row_u0 = index.mul %ln4, %k_wstride : index + %rfs2_2_1_base_u0 = index.add %rfs2_2_1_row_u0, %blk_k16_u0 : index + %rfs2_2_1_idx_u0 = index.assume %rfs2_2_1_base_u0 [range(%rfs2_2_1_base_u0, 0, 5112)] : index + %rfs2_2_1_raw_u0 = vector.load %wl_flat[%rfs2_2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_2_1_words_u0 = vector.bitcast %rfs2_2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_2_1_u0 = vector.fragment %rfs2_2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_2_u0 = vector.mma %lf2_0_u0, %rfs2_2_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_2_u0 = vector.mma %lf2_1_u0, %rfs2_2_1_u0, %i0_2_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff2_2_u0 = vector.sitofp %i1_2_2_u0 : vector<8xi32> to vector<8xf32> + %sv2_2_u0 = vector.mulf %asv2_u0, %wsv4_u0 : vector<8xf32> + %fn2_2_u0 = vector.fmaf %ff2_2_u0, %sv2_2_u0, %fc2_2 : vector<8xf32> + scf.schedule.fence + + %rfs2_3_0_row_u0 = index.mul %ln5, %k_wstride : index + %rfs2_3_0_base_u0 = index.add %rfs2_3_0_row_u0, %blk_k_u0 : index + %rfs2_3_0_idx_u0 = index.assume %rfs2_3_0_base_u0 [range(%rfs2_3_0_base_u0, 0, 5112)] : index + %rfs2_3_0_raw_u0 = vector.load %wl_flat[%rfs2_3_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_3_0_words_u0 = vector.bitcast %rfs2_3_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_3_0_u0 = vector.fragment %rfs2_3_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_3_1_row_u0 = index.mul %ln5, %k_wstride : index + %rfs2_3_1_base_u0 = index.add %rfs2_3_1_row_u0, %blk_k16_u0 : index + %rfs2_3_1_idx_u0 = index.assume %rfs2_3_1_base_u0 [range(%rfs2_3_1_base_u0, 0, 5112)] : index + %rfs2_3_1_raw_u0 = vector.load %wl_flat[%rfs2_3_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs2_3_1_words_u0 = vector.bitcast %rfs2_3_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs2_3_1_u0 = vector.fragment %rfs2_3_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_3_u0 = vector.mma %lf2_0_u0, %rfs2_3_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_3_u0 = vector.mma %lf2_1_u0, %rfs2_3_1_u0, %i0_2_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff2_3_u0 = vector.sitofp %i1_2_3_u0 : vector<8xi32> to vector<8xf32> + %sv2_3_u0 = vector.mulf %asv2_u0, %wsv5_u0 : vector<8xf32> + %fn2_3_u0 = vector.fmaf %ff2_3_u0, %sv2_3_u0, %fc2_3 : vector<8xf32> + + scf.schedule.fence + %rfs3_0_0_row_u0 = index.mul %ln2, %k_wstride : index + %rfs3_0_0_base_u0 = index.add %rfs3_0_0_row_u0, %blk_k_u0 : index + %rfs3_0_0_idx_u0 = index.assume %rfs3_0_0_base_u0 [range(%rfs3_0_0_base_u0, 0, 5112)] : index + %rfs3_0_0_raw_u0 = vector.load %wl_flat[%rfs3_0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_0_0_words_u0 = vector.bitcast %rfs3_0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_0_0_u0 = vector.fragment %rfs3_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_0_1_row_u0 = index.mul %ln2, %k_wstride : index + %rfs3_0_1_base_u0 = index.add %rfs3_0_1_row_u0, %blk_k16_u0 : index + %rfs3_0_1_idx_u0 = index.assume %rfs3_0_1_base_u0 [range(%rfs3_0_1_base_u0, 0, 5112)] : index + %rfs3_0_1_raw_u0 = vector.load %wl_flat[%rfs3_0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_0_1_words_u0 = vector.bitcast %rfs3_0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_0_1_u0 = vector.fragment %rfs3_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_0_u0 = vector.mma %lf2_0_u0, %rfs3_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_0_u0 = vector.mma %lf2_1_u0, %rfs3_0_1_u0, %i0_3_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff3_0_u0 = vector.sitofp %i1_3_0_u0 : vector<8xi32> to vector<8xf32> + %sv3_0_u0 = vector.mulf %asv2_u0, %wsv2_u0 : vector<8xf32> + %fn3_0_u0 = vector.fmaf %ff3_0_u0, %sv3_0_u0, %fc3_0 : vector<8xf32> + scf.schedule.fence + + %rfs3_1_0_row_u0 = index.mul %ln3, %k_wstride : index + %rfs3_1_0_base_u0 = index.add %rfs3_1_0_row_u0, %blk_k_u0 : index + %rfs3_1_0_idx_u0 = index.assume %rfs3_1_0_base_u0 [range(%rfs3_1_0_base_u0, 0, 5112)] : index + %rfs3_1_0_raw_u0 = vector.load %wl_flat[%rfs3_1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_1_0_words_u0 = vector.bitcast %rfs3_1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_1_0_u0 = vector.fragment %rfs3_1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_1_1_row_u0 = index.mul %ln3, %k_wstride : index + %rfs3_1_1_base_u0 = index.add %rfs3_1_1_row_u0, %blk_k16_u0 : index + %rfs3_1_1_idx_u0 = index.assume %rfs3_1_1_base_u0 [range(%rfs3_1_1_base_u0, 0, 5112)] : index + %rfs3_1_1_raw_u0 = vector.load %wl_flat[%rfs3_1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_1_1_words_u0 = vector.bitcast %rfs3_1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_1_1_u0 = vector.fragment %rfs3_1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_1_u0 = vector.mma %lf2_0_u0, %rfs3_1_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_1_u0 = vector.mma %lf2_1_u0, %rfs3_1_1_u0, %i0_3_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff3_1_u0 = vector.sitofp %i1_3_1_u0 : vector<8xi32> to vector<8xf32> + %sv3_1_u0 = vector.mulf %asv2_u0, %wsv3_u0 : vector<8xf32> + %fn3_1_u0 = vector.fmaf %ff3_1_u0, %sv3_1_u0, %fc3_1 : vector<8xf32> + scf.schedule.fence + + %rfs3_2_0_row_u0 = index.mul %ln6, %k_wstride : index + %rfs3_2_0_base_u0 = index.add %rfs3_2_0_row_u0, %blk_k_u0 : index + %rfs3_2_0_idx_u0 = index.assume %rfs3_2_0_base_u0 [range(%rfs3_2_0_base_u0, 0, 5112)] : index + %rfs3_2_0_raw_u0 = vector.load %wl_flat[%rfs3_2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_2_0_words_u0 = vector.bitcast %rfs3_2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_2_0_u0 = vector.fragment %rfs3_2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_2_1_row_u0 = index.mul %ln6, %k_wstride : index + %rfs3_2_1_base_u0 = index.add %rfs3_2_1_row_u0, %blk_k16_u0 : index + %rfs3_2_1_idx_u0 = index.assume %rfs3_2_1_base_u0 [range(%rfs3_2_1_base_u0, 0, 5112)] : index + %rfs3_2_1_raw_u0 = vector.load %wl_flat[%rfs3_2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_2_1_words_u0 = vector.bitcast %rfs3_2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_2_1_u0 = vector.fragment %rfs3_2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_2_u0 = vector.mma %lf2_0_u0, %rfs3_2_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_2_u0 = vector.mma %lf2_1_u0, %rfs3_2_1_u0, %i0_3_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff3_2_u0 = vector.sitofp %i1_3_2_u0 : vector<8xi32> to vector<8xf32> + %sv3_2_u0 = vector.mulf %asv2_u0, %wsv6_u0 : vector<8xf32> + %fn3_2_u0 = vector.fmaf %ff3_2_u0, %sv3_2_u0, %fc3_2 : vector<8xf32> + scf.schedule.fence + + %rfs3_3_0_row_u0 = index.mul %ln7, %k_wstride : index + %rfs3_3_0_base_u0 = index.add %rfs3_3_0_row_u0, %blk_k_u0 : index + %rfs3_3_0_idx_u0 = index.assume %rfs3_3_0_base_u0 [range(%rfs3_3_0_base_u0, 0, 5112)] : index + %rfs3_3_0_raw_u0 = vector.load %wl_flat[%rfs3_3_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_3_0_words_u0 = vector.bitcast %rfs3_3_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_3_0_u0 = vector.fragment %rfs3_3_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_3_1_row_u0 = index.mul %ln7, %k_wstride : index + %rfs3_3_1_base_u0 = index.add %rfs3_3_1_row_u0, %blk_k16_u0 : index + %rfs3_3_1_idx_u0 = index.assume %rfs3_3_1_base_u0 [range(%rfs3_3_1_base_u0, 0, 5112)] : index + %rfs3_3_1_raw_u0 = vector.load %wl_flat[%rfs3_3_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs3_3_1_words_u0 = vector.bitcast %rfs3_3_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs3_3_1_u0 = vector.fragment %rfs3_3_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_3_u0 = vector.mma %lf2_0_u0, %rfs3_3_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_3_u0 = vector.mma %lf2_1_u0, %rfs3_3_1_u0, %i0_3_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff3_3_u0 = vector.sitofp %i1_3_3_u0 : vector<8xi32> to vector<8xf32> + %sv3_3_u0 = vector.mulf %asv2_u0, %wsv7_u0 : vector<8xf32> + %fn3_3_u0 = vector.fmaf %ff3_3_u0, %sv3_3_u0, %fc3_3 : vector<8xf32> + + scf.schedule.fence + %blk_k_u1 = index.mul %c1, %c16 : index + %blk_k16_u1 = index.add %blk_k_u1, %c8 : index + %wsi0_2_u1 = index.add %wsb0, %c1 : index + %wsi0_u1 = index.assume %wsi0_2_u1 [range(%wsi0_2_u1, 1, 255)] : index + %wsc0_u1 = view.load %wsl_view[%wsi0_u1] : view<256xf32> -> f32 + %wsv0_u1 = vector.splat %wsc0_u1 : vector<8xf32> + %wsi1_2_u1 = index.add %wsb1, %c1 : index + %wsi1_u1 = index.assume %wsi1_2_u1 [range(%wsi1_2_u1, 1, 255)] : index + %wsc1_u1 = view.load %wsl_view[%wsi1_u1] : view<256xf32> -> f32 + %wsv1_u1 = vector.splat %wsc1_u1 : vector<8xf32> + %wsi4_2_u1 = index.add %wsb4, %c1 : index + %wsi4_u1 = index.assume %wsi4_2_u1 [range(%wsi4_2_u1, 1, 255)] : index + %wsc4_u1 = view.load %wsl_view[%wsi4_u1] : view<256xf32> -> f32 + %wsv4_u1 = vector.splat %wsc4_u1 : vector<8xf32> + %wsi5_2_u1 = index.add %wsb5, %c1 : index + %wsi5_u1 = index.assume %wsi5_2_u1 [range(%wsi5_2_u1, 1, 255)] : index + %wsc5_u1 = view.load %wsl_view[%wsi5_u1] : view<256xf32> -> f32 + %wsv5_u1 = vector.splat %wsc5_u1 : vector<8xf32> + %lf0_0_lane_row_u1 = index.add %lm0, %lane_lo : index + %lf0_0_row_u1 = index.mul %lf0_0_lane_row_u1, %k_astride : index + %lf0_0_base_u1 = index.add %lf0_0_row_u1, %blk_k_u1 : index + %lf0_0_idx_u1 = index.assume %lf0_0_base_u1 [range(%lf0_0_base_u1, 0, 5112)] : index + %lf0_0_raw_u1 = vector.load %al_flat[%lf0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf0_0_words_u1 = vector.bitcast %lf0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %lf0_0_u1 = vector.fragment %lf0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf0_1_lane_row_u1 = index.add %lm0, %lane_lo : index + %lf0_1_row_u1 = index.mul %lf0_1_lane_row_u1, %k_astride : index + %lf0_1_base_u1 = index.add %lf0_1_row_u1, %blk_k16_u1 : index + %lf0_1_idx_u1 = index.assume %lf0_1_base_u1 [range(%lf0_1_base_u1, 0, 5112)] : index + %lf0_1_raw_u1 = vector.load %al_flat[%lf0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf0_1_words_u1 = vector.bitcast %lf0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %lf0_1_u1 = vector.fragment %lf0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_0_row_u1 = index.mul %ln0, %k_wstride : index + %rfs0_0_0_base_u1 = index.add %rfs0_0_0_row_u1, %blk_k_u1 : index + %rfs0_0_0_idx_u1 = index.assume %rfs0_0_0_base_u1 [range(%rfs0_0_0_base_u1, 0, 5112)] : index + %rfs0_0_0_raw_u1 = vector.load %wl_flat[%rfs0_0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_0_0_words_u1 = vector.bitcast %rfs0_0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_0_0_u1 = vector.fragment %rfs0_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_1_row_u1 = index.mul %ln0, %k_wstride : index + %rfs0_0_1_base_u1 = index.add %rfs0_0_1_row_u1, %blk_k16_u1 : index + %rfs0_0_1_idx_u1 = index.assume %rfs0_0_1_base_u1 [range(%rfs0_0_1_base_u1, 0, 5112)] : index + %rfs0_0_1_raw_u1 = vector.load %wl_flat[%rfs0_0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_0_1_words_u1 = vector.bitcast %rfs0_0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_0_1_u1 = vector.fragment %rfs0_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_0_u1 = vector.mma %lf0_0_u1, %rfs0_0_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_0_u1 = vector.mma %lf0_1_u1, %rfs0_0_1_u1, %i0_0_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff0_0_u1 = vector.sitofp %i1_0_0_u1 : vector<8xi32> to vector<8xf32> + %sv0_0_u1 = vector.mulf %asv0_u0, %wsv0_u1 : vector<8xf32> + %fn0_0 = vector.fmaf %ff0_0_u1, %sv0_0_u1, %fn0_0_u0 : vector<8xf32> + scf.schedule.fence + %rfs0_1_0_row_u1 = index.mul %ln1, %k_wstride : index + %rfs0_1_0_base_u1 = index.add %rfs0_1_0_row_u1, %blk_k_u1 : index + %rfs0_1_0_idx_u1 = index.assume %rfs0_1_0_base_u1 [range(%rfs0_1_0_base_u1, 0, 5112)] : index + %rfs0_1_0_raw_u1 = vector.load %wl_flat[%rfs0_1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_1_0_words_u1 = vector.bitcast %rfs0_1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_1_0_u1 = vector.fragment %rfs0_1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_1_1_row_u1 = index.mul %ln1, %k_wstride : index + %rfs0_1_1_base_u1 = index.add %rfs0_1_1_row_u1, %blk_k16_u1 : index + %rfs0_1_1_idx_u1 = index.assume %rfs0_1_1_base_u1 [range(%rfs0_1_1_base_u1, 0, 5112)] : index + %rfs0_1_1_raw_u1 = vector.load %wl_flat[%rfs0_1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_1_1_words_u1 = vector.bitcast %rfs0_1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_1_1_u1 = vector.fragment %rfs0_1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_1_u1 = vector.mma %lf0_0_u1, %rfs0_1_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_1_u1 = vector.mma %lf0_1_u1, %rfs0_1_1_u1, %i0_0_1_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff0_1_u1 = vector.sitofp %i1_0_1_u1 : vector<8xi32> to vector<8xf32> + %sv0_1_u1 = vector.mulf %asv0_u0, %wsv1_u1 : vector<8xf32> + %fn0_1 = vector.fmaf %ff0_1_u1, %sv0_1_u1, %fn0_1_u0 : vector<8xf32> + scf.schedule.fence + %rfs0_2_0_row_u1 = index.mul %ln4, %k_wstride : index + %rfs0_2_0_base_u1 = index.add %rfs0_2_0_row_u1, %blk_k_u1 : index + %rfs0_2_0_idx_u1 = index.assume %rfs0_2_0_base_u1 [range(%rfs0_2_0_base_u1, 0, 5112)] : index + %rfs0_2_0_raw_u1 = vector.load %wl_flat[%rfs0_2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_2_0_words_u1 = vector.bitcast %rfs0_2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_2_0_u1 = vector.fragment %rfs0_2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_2_1_row_u1 = index.mul %ln4, %k_wstride : index + %rfs0_2_1_base_u1 = index.add %rfs0_2_1_row_u1, %blk_k16_u1 : index + %rfs0_2_1_idx_u1 = index.assume %rfs0_2_1_base_u1 [range(%rfs0_2_1_base_u1, 0, 5112)] : index + %rfs0_2_1_raw_u1 = vector.load %wl_flat[%rfs0_2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_2_1_words_u1 = vector.bitcast %rfs0_2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_2_1_u1 = vector.fragment %rfs0_2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_2_u1 = vector.mma %lf0_0_u1, %rfs0_2_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_2_u1 = vector.mma %lf0_1_u1, %rfs0_2_1_u1, %i0_0_2_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff0_2_u1 = vector.sitofp %i1_0_2_u1 : vector<8xi32> to vector<8xf32> + %sv0_2_u1 = vector.mulf %asv0_u0, %wsv4_u1 : vector<8xf32> + %fn0_2 = vector.fmaf %ff0_2_u1, %sv0_2_u1, %fn0_2_u0 : vector<8xf32> + scf.schedule.fence + %rfs0_3_0_row_u1 = index.mul %ln5, %k_wstride : index + %rfs0_3_0_base_u1 = index.add %rfs0_3_0_row_u1, %blk_k_u1 : index + %rfs0_3_0_idx_u1 = index.assume %rfs0_3_0_base_u1 [range(%rfs0_3_0_base_u1, 0, 5112)] : index + %rfs0_3_0_raw_u1 = vector.load %wl_flat[%rfs0_3_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_3_0_words_u1 = vector.bitcast %rfs0_3_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_3_0_u1 = vector.fragment %rfs0_3_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_3_1_row_u1 = index.mul %ln5, %k_wstride : index + %rfs0_3_1_base_u1 = index.add %rfs0_3_1_row_u1, %blk_k16_u1 : index + %rfs0_3_1_idx_u1 = index.assume %rfs0_3_1_base_u1 [range(%rfs0_3_1_base_u1, 0, 5112)] : index + %rfs0_3_1_raw_u1 = vector.load %wl_flat[%rfs0_3_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_3_1_words_u1 = vector.bitcast %rfs0_3_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_3_1_u1 = vector.fragment %rfs0_3_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_3_u1 = vector.mma %lf0_0_u1, %rfs0_3_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_3_u1 = vector.mma %lf0_1_u1, %rfs0_3_1_u1, %i0_0_3_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff0_3_u1 = vector.sitofp %i1_0_3_u1 : vector<8xi32> to vector<8xf32> + %sv0_3_u1 = vector.mulf %asv0_u0, %wsv5_u1 : vector<8xf32> + %fn0_3 = vector.fmaf %ff0_3_u1, %sv0_3_u1, %fn0_3_u0 : vector<8xf32> + scf.schedule.fence + %wsi2_2_u1 = index.add %wsb2, %c1 : index + %wsi2_u1 = index.assume %wsi2_2_u1 [range(%wsi2_2_u1, 1, 255)] : index + %wsc2_u1 = view.load %wsl_view[%wsi2_u1] : view<256xf32> -> f32 + %wsv2_u1 = vector.splat %wsc2_u1 : vector<8xf32> + %wsi3_2_u1 = index.add %wsb3, %c1 : index + %wsi3_u1 = index.assume %wsi3_2_u1 [range(%wsi3_2_u1, 1, 255)] : index + %wsc3_u1 = view.load %wsl_view[%wsi3_u1] : view<256xf32> -> f32 + %wsv3_u1 = vector.splat %wsc3_u1 : vector<8xf32> + %wsi6_2_u1 = index.add %wsb6, %c1 : index + %wsi6_u1 = index.assume %wsi6_2_u1 [range(%wsi6_2_u1, 1, 255)] : index + %wsc6_u1 = view.load %wsl_view[%wsi6_u1] : view<256xf32> -> f32 + %wsv6_u1 = vector.splat %wsc6_u1 : vector<8xf32> + %wsi7_2_u1 = index.add %wsb7, %c1 : index + %wsi7_u1 = index.assume %wsi7_2_u1 [range(%wsi7_2_u1, 1, 255)] : index + %wsc7_u1 = view.load %wsl_view[%wsi7_u1] : view<256xf32> -> f32 + %wsv7_u1 = vector.splat %wsc7_u1 : vector<8xf32> + %rfs1_0_0_row_u1 = index.mul %ln2, %k_wstride : index + %rfs1_0_0_base_u1 = index.add %rfs1_0_0_row_u1, %blk_k_u1 : index + %rfs1_0_0_idx_u1 = index.assume %rfs1_0_0_base_u1 [range(%rfs1_0_0_base_u1, 0, 5112)] : index + %rfs1_0_0_raw_u1 = vector.load %wl_flat[%rfs1_0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_0_0_words_u1 = vector.bitcast %rfs1_0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_0_0_u1 = vector.fragment %rfs1_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_0_1_row_u1 = index.mul %ln2, %k_wstride : index + %rfs1_0_1_base_u1 = index.add %rfs1_0_1_row_u1, %blk_k16_u1 : index + %rfs1_0_1_idx_u1 = index.assume %rfs1_0_1_base_u1 [range(%rfs1_0_1_base_u1, 0, 5112)] : index + %rfs1_0_1_raw_u1 = vector.load %wl_flat[%rfs1_0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_0_1_words_u1 = vector.bitcast %rfs1_0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_0_1_u1 = vector.fragment %rfs1_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_0_u1 = vector.mma %lf0_0_u1, %rfs1_0_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_0_u1 = vector.mma %lf0_1_u1, %rfs1_0_1_u1, %i0_1_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff1_0_u1 = vector.sitofp %i1_1_0_u1 : vector<8xi32> to vector<8xf32> + %sv1_0_u1 = vector.mulf %asv0_u0, %wsv2_u1 : vector<8xf32> + %fn1_0 = vector.fmaf %ff1_0_u1, %sv1_0_u1, %fn1_0_u0 : vector<8xf32> + scf.schedule.fence + %rfs1_1_0_row_u1 = index.mul %ln3, %k_wstride : index + %rfs1_1_0_base_u1 = index.add %rfs1_1_0_row_u1, %blk_k_u1 : index + %rfs1_1_0_idx_u1 = index.assume %rfs1_1_0_base_u1 [range(%rfs1_1_0_base_u1, 0, 5112)] : index + %rfs1_1_0_raw_u1 = vector.load %wl_flat[%rfs1_1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_1_0_words_u1 = vector.bitcast %rfs1_1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_1_0_u1 = vector.fragment %rfs1_1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_1_1_row_u1 = index.mul %ln3, %k_wstride : index + %rfs1_1_1_base_u1 = index.add %rfs1_1_1_row_u1, %blk_k16_u1 : index + %rfs1_1_1_idx_u1 = index.assume %rfs1_1_1_base_u1 [range(%rfs1_1_1_base_u1, 0, 5112)] : index + %rfs1_1_1_raw_u1 = vector.load %wl_flat[%rfs1_1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_1_1_words_u1 = vector.bitcast %rfs1_1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_1_1_u1 = vector.fragment %rfs1_1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_1_u1 = vector.mma %lf0_0_u1, %rfs1_1_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_1_u1 = vector.mma %lf0_1_u1, %rfs1_1_1_u1, %i0_1_1_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff1_1_u1 = vector.sitofp %i1_1_1_u1 : vector<8xi32> to vector<8xf32> + %sv1_1_u1 = vector.mulf %asv0_u0, %wsv3_u1 : vector<8xf32> + %fn1_1 = vector.fmaf %ff1_1_u1, %sv1_1_u1, %fn1_1_u0 : vector<8xf32> + scf.schedule.fence + %rfs1_2_0_row_u1 = index.mul %ln6, %k_wstride : index + %rfs1_2_0_base_u1 = index.add %rfs1_2_0_row_u1, %blk_k_u1 : index + %rfs1_2_0_idx_u1 = index.assume %rfs1_2_0_base_u1 [range(%rfs1_2_0_base_u1, 0, 5112)] : index + %rfs1_2_0_raw_u1 = vector.load %wl_flat[%rfs1_2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_2_0_words_u1 = vector.bitcast %rfs1_2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_2_0_u1 = vector.fragment %rfs1_2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_2_1_row_u1 = index.mul %ln6, %k_wstride : index + %rfs1_2_1_base_u1 = index.add %rfs1_2_1_row_u1, %blk_k16_u1 : index + %rfs1_2_1_idx_u1 = index.assume %rfs1_2_1_base_u1 [range(%rfs1_2_1_base_u1, 0, 5112)] : index + %rfs1_2_1_raw_u1 = vector.load %wl_flat[%rfs1_2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_2_1_words_u1 = vector.bitcast %rfs1_2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_2_1_u1 = vector.fragment %rfs1_2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_2_u1 = vector.mma %lf0_0_u1, %rfs1_2_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_2_u1 = vector.mma %lf0_1_u1, %rfs1_2_1_u1, %i0_1_2_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff1_2_u1 = vector.sitofp %i1_1_2_u1 : vector<8xi32> to vector<8xf32> + %sv1_2_u1 = vector.mulf %asv0_u0, %wsv6_u1 : vector<8xf32> + %fn1_2 = vector.fmaf %ff1_2_u1, %sv1_2_u1, %fn1_2_u0 : vector<8xf32> + scf.schedule.fence + %rfs1_3_0_row_u1 = index.mul %ln7, %k_wstride : index + %rfs1_3_0_base_u1 = index.add %rfs1_3_0_row_u1, %blk_k_u1 : index + %rfs1_3_0_idx_u1 = index.assume %rfs1_3_0_base_u1 [range(%rfs1_3_0_base_u1, 0, 5112)] : index + %rfs1_3_0_raw_u1 = vector.load %wl_flat[%rfs1_3_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_3_0_words_u1 = vector.bitcast %rfs1_3_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_3_0_u1 = vector.fragment %rfs1_3_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_3_1_row_u1 = index.mul %ln7, %k_wstride : index + %rfs1_3_1_base_u1 = index.add %rfs1_3_1_row_u1, %blk_k16_u1 : index + %rfs1_3_1_idx_u1 = index.assume %rfs1_3_1_base_u1 [range(%rfs1_3_1_base_u1, 0, 5112)] : index + %rfs1_3_1_raw_u1 = vector.load %wl_flat[%rfs1_3_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_3_1_words_u1 = vector.bitcast %rfs1_3_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_3_1_u1 = vector.fragment %rfs1_3_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_3_u1 = vector.mma %lf0_0_u1, %rfs1_3_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_3_u1 = vector.mma %lf0_1_u1, %rfs1_3_1_u1, %i0_1_3_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff1_3_u1 = vector.sitofp %i1_1_3_u1 : vector<8xi32> to vector<8xf32> + %sv1_3_u1 = vector.mulf %asv0_u0, %wsv7_u1 : vector<8xf32> + %fn1_3 = vector.fmaf %ff1_3_u1, %sv1_3_u1, %fn1_3_u0 : vector<8xf32> + scf.schedule.fence + %lf2_0_lane_row_u1 = index.add %lm2, %lane_lo : index + %lf2_0_row_u1 = index.mul %lf2_0_lane_row_u1, %k_astride : index + %lf2_0_base_u1 = index.add %lf2_0_row_u1, %blk_k_u1 : index + %lf2_0_idx_u1 = index.assume %lf2_0_base_u1 [range(%lf2_0_base_u1, 0, 5112)] : index + %lf2_0_raw_u1 = vector.load %al_flat[%lf2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf2_0_words_u1 = vector.bitcast %lf2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %lf2_0_u1 = vector.fragment %lf2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf2_1_lane_row_u1 = index.add %lm2, %lane_lo : index + %lf2_1_row_u1 = index.mul %lf2_1_lane_row_u1, %k_astride : index + %lf2_1_base_u1 = index.add %lf2_1_row_u1, %blk_k16_u1 : index + %lf2_1_idx_u1 = index.assume %lf2_1_base_u1 [range(%lf2_1_base_u1, 0, 5112)] : index + %lf2_1_raw_u1 = vector.load %al_flat[%lf2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf2_1_words_u1 = vector.bitcast %lf2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %lf2_1_u1 = vector.fragment %lf2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_0_0_row_u1 = index.mul %ln0, %k_wstride : index + %rfs2_0_0_base_u1 = index.add %rfs2_0_0_row_u1, %blk_k_u1 : index + %rfs2_0_0_idx_u1 = index.assume %rfs2_0_0_base_u1 [range(%rfs2_0_0_base_u1, 0, 5112)] : index + %rfs2_0_0_raw_u1 = vector.load %wl_flat[%rfs2_0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_0_0_words_u1 = vector.bitcast %rfs2_0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_0_0_u1 = vector.fragment %rfs2_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_0_1_row_u1 = index.mul %ln0, %k_wstride : index + %rfs2_0_1_base_u1 = index.add %rfs2_0_1_row_u1, %blk_k16_u1 : index + %rfs2_0_1_idx_u1 = index.assume %rfs2_0_1_base_u1 [range(%rfs2_0_1_base_u1, 0, 5112)] : index + %rfs2_0_1_raw_u1 = vector.load %wl_flat[%rfs2_0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_0_1_words_u1 = vector.bitcast %rfs2_0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_0_1_u1 = vector.fragment %rfs2_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_0_u1 = vector.mma %lf2_0_u1, %rfs2_0_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_0_u1 = vector.mma %lf2_1_u1, %rfs2_0_1_u1, %i0_2_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff2_0_u1 = vector.sitofp %i1_2_0_u1 : vector<8xi32> to vector<8xf32> + %sv2_0_u1 = vector.mulf %asv2_u0, %wsv0_u1 : vector<8xf32> + %fn2_0 = vector.fmaf %ff2_0_u1, %sv2_0_u1, %fn2_0_u0 : vector<8xf32> + scf.schedule.fence + %rfs2_1_0_row_u1 = index.mul %ln1, %k_wstride : index + %rfs2_1_0_base_u1 = index.add %rfs2_1_0_row_u1, %blk_k_u1 : index + %rfs2_1_0_idx_u1 = index.assume %rfs2_1_0_base_u1 [range(%rfs2_1_0_base_u1, 0, 5112)] : index + %rfs2_1_0_raw_u1 = vector.load %wl_flat[%rfs2_1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_1_0_words_u1 = vector.bitcast %rfs2_1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_1_0_u1 = vector.fragment %rfs2_1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_1_1_row_u1 = index.mul %ln1, %k_wstride : index + %rfs2_1_1_base_u1 = index.add %rfs2_1_1_row_u1, %blk_k16_u1 : index + %rfs2_1_1_idx_u1 = index.assume %rfs2_1_1_base_u1 [range(%rfs2_1_1_base_u1, 0, 5112)] : index + %rfs2_1_1_raw_u1 = vector.load %wl_flat[%rfs2_1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_1_1_words_u1 = vector.bitcast %rfs2_1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_1_1_u1 = vector.fragment %rfs2_1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_1_u1 = vector.mma %lf2_0_u1, %rfs2_1_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_1_u1 = vector.mma %lf2_1_u1, %rfs2_1_1_u1, %i0_2_1_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff2_1_u1 = vector.sitofp %i1_2_1_u1 : vector<8xi32> to vector<8xf32> + %sv2_1_u1 = vector.mulf %asv2_u0, %wsv1_u1 : vector<8xf32> + %fn2_1 = vector.fmaf %ff2_1_u1, %sv2_1_u1, %fn2_1_u0 : vector<8xf32> + scf.schedule.fence + %rfs2_2_0_row_u1 = index.mul %ln4, %k_wstride : index + %rfs2_2_0_base_u1 = index.add %rfs2_2_0_row_u1, %blk_k_u1 : index + %rfs2_2_0_idx_u1 = index.assume %rfs2_2_0_base_u1 [range(%rfs2_2_0_base_u1, 0, 5112)] : index + %rfs2_2_0_raw_u1 = vector.load %wl_flat[%rfs2_2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_2_0_words_u1 = vector.bitcast %rfs2_2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_2_0_u1 = vector.fragment %rfs2_2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_2_1_row_u1 = index.mul %ln4, %k_wstride : index + %rfs2_2_1_base_u1 = index.add %rfs2_2_1_row_u1, %blk_k16_u1 : index + %rfs2_2_1_idx_u1 = index.assume %rfs2_2_1_base_u1 [range(%rfs2_2_1_base_u1, 0, 5112)] : index + %rfs2_2_1_raw_u1 = vector.load %wl_flat[%rfs2_2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_2_1_words_u1 = vector.bitcast %rfs2_2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_2_1_u1 = vector.fragment %rfs2_2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_2_u1 = vector.mma %lf2_0_u1, %rfs2_2_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_2_u1 = vector.mma %lf2_1_u1, %rfs2_2_1_u1, %i0_2_2_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff2_2_u1 = vector.sitofp %i1_2_2_u1 : vector<8xi32> to vector<8xf32> + %sv2_2_u1 = vector.mulf %asv2_u0, %wsv4_u1 : vector<8xf32> + %fn2_2 = vector.fmaf %ff2_2_u1, %sv2_2_u1, %fn2_2_u0 : vector<8xf32> + scf.schedule.fence + %rfs2_3_0_row_u1 = index.mul %ln5, %k_wstride : index + %rfs2_3_0_base_u1 = index.add %rfs2_3_0_row_u1, %blk_k_u1 : index + %rfs2_3_0_idx_u1 = index.assume %rfs2_3_0_base_u1 [range(%rfs2_3_0_base_u1, 0, 5112)] : index + %rfs2_3_0_raw_u1 = vector.load %wl_flat[%rfs2_3_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_3_0_words_u1 = vector.bitcast %rfs2_3_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_3_0_u1 = vector.fragment %rfs2_3_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs2_3_1_row_u1 = index.mul %ln5, %k_wstride : index + %rfs2_3_1_base_u1 = index.add %rfs2_3_1_row_u1, %blk_k16_u1 : index + %rfs2_3_1_idx_u1 = index.assume %rfs2_3_1_base_u1 [range(%rfs2_3_1_base_u1, 0, 5112)] : index + %rfs2_3_1_raw_u1 = vector.load %wl_flat[%rfs2_3_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs2_3_1_words_u1 = vector.bitcast %rfs2_3_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs2_3_1_u1 = vector.fragment %rfs2_3_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_2_3_u1 = vector.mma %lf2_0_u1, %rfs2_3_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_2_3_u1 = vector.mma %lf2_1_u1, %rfs2_3_1_u1, %i0_2_3_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff2_3_u1 = vector.sitofp %i1_2_3_u1 : vector<8xi32> to vector<8xf32> + %sv2_3_u1 = vector.mulf %asv2_u0, %wsv5_u1 : vector<8xf32> + %fn2_3 = vector.fmaf %ff2_3_u1, %sv2_3_u1, %fn2_3_u0 : vector<8xf32> + scf.schedule.fence + %rfs3_0_0_row_u1 = index.mul %ln2, %k_wstride : index + %rfs3_0_0_base_u1 = index.add %rfs3_0_0_row_u1, %blk_k_u1 : index + %rfs3_0_0_idx_u1 = index.assume %rfs3_0_0_base_u1 [range(%rfs3_0_0_base_u1, 0, 5112)] : index + %rfs3_0_0_raw_u1 = vector.load %wl_flat[%rfs3_0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_0_0_words_u1 = vector.bitcast %rfs3_0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_0_0_u1 = vector.fragment %rfs3_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_0_1_row_u1 = index.mul %ln2, %k_wstride : index + %rfs3_0_1_base_u1 = index.add %rfs3_0_1_row_u1, %blk_k16_u1 : index + %rfs3_0_1_idx_u1 = index.assume %rfs3_0_1_base_u1 [range(%rfs3_0_1_base_u1, 0, 5112)] : index + %rfs3_0_1_raw_u1 = vector.load %wl_flat[%rfs3_0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_0_1_words_u1 = vector.bitcast %rfs3_0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_0_1_u1 = vector.fragment %rfs3_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_0_u1 = vector.mma %lf2_0_u1, %rfs3_0_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_0_u1 = vector.mma %lf2_1_u1, %rfs3_0_1_u1, %i0_3_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff3_0_u1 = vector.sitofp %i1_3_0_u1 : vector<8xi32> to vector<8xf32> + %sv3_0_u1 = vector.mulf %asv2_u0, %wsv2_u1 : vector<8xf32> + %fn3_0 = vector.fmaf %ff3_0_u1, %sv3_0_u1, %fn3_0_u0 : vector<8xf32> + scf.schedule.fence + %rfs3_1_0_row_u1 = index.mul %ln3, %k_wstride : index + %rfs3_1_0_base_u1 = index.add %rfs3_1_0_row_u1, %blk_k_u1 : index + %rfs3_1_0_idx_u1 = index.assume %rfs3_1_0_base_u1 [range(%rfs3_1_0_base_u1, 0, 5112)] : index + %rfs3_1_0_raw_u1 = vector.load %wl_flat[%rfs3_1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_1_0_words_u1 = vector.bitcast %rfs3_1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_1_0_u1 = vector.fragment %rfs3_1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_1_1_row_u1 = index.mul %ln3, %k_wstride : index + %rfs3_1_1_base_u1 = index.add %rfs3_1_1_row_u1, %blk_k16_u1 : index + %rfs3_1_1_idx_u1 = index.assume %rfs3_1_1_base_u1 [range(%rfs3_1_1_base_u1, 0, 5112)] : index + %rfs3_1_1_raw_u1 = vector.load %wl_flat[%rfs3_1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_1_1_words_u1 = vector.bitcast %rfs3_1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_1_1_u1 = vector.fragment %rfs3_1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_1_u1 = vector.mma %lf2_0_u1, %rfs3_1_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_1_u1 = vector.mma %lf2_1_u1, %rfs3_1_1_u1, %i0_3_1_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff3_1_u1 = vector.sitofp %i1_3_1_u1 : vector<8xi32> to vector<8xf32> + %sv3_1_u1 = vector.mulf %asv2_u0, %wsv3_u1 : vector<8xf32> + %fn3_1 = vector.fmaf %ff3_1_u1, %sv3_1_u1, %fn3_1_u0 : vector<8xf32> + scf.schedule.fence + %rfs3_2_0_row_u1 = index.mul %ln6, %k_wstride : index + %rfs3_2_0_base_u1 = index.add %rfs3_2_0_row_u1, %blk_k_u1 : index + %rfs3_2_0_idx_u1 = index.assume %rfs3_2_0_base_u1 [range(%rfs3_2_0_base_u1, 0, 5112)] : index + %rfs3_2_0_raw_u1 = vector.load %wl_flat[%rfs3_2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_2_0_words_u1 = vector.bitcast %rfs3_2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_2_0_u1 = vector.fragment %rfs3_2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_2_1_row_u1 = index.mul %ln6, %k_wstride : index + %rfs3_2_1_base_u1 = index.add %rfs3_2_1_row_u1, %blk_k16_u1 : index + %rfs3_2_1_idx_u1 = index.assume %rfs3_2_1_base_u1 [range(%rfs3_2_1_base_u1, 0, 5112)] : index + %rfs3_2_1_raw_u1 = vector.load %wl_flat[%rfs3_2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_2_1_words_u1 = vector.bitcast %rfs3_2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_2_1_u1 = vector.fragment %rfs3_2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_2_u1 = vector.mma %lf2_0_u1, %rfs3_2_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_2_u1 = vector.mma %lf2_1_u1, %rfs3_2_1_u1, %i0_3_2_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff3_2_u1 = vector.sitofp %i1_3_2_u1 : vector<8xi32> to vector<8xf32> + %sv3_2_u1 = vector.mulf %asv2_u0, %wsv6_u1 : vector<8xf32> + %fn3_2 = vector.fmaf %ff3_2_u1, %sv3_2_u1, %fn3_2_u0 : vector<8xf32> + scf.schedule.fence + %rfs3_3_0_row_u1 = index.mul %ln7, %k_wstride : index + %rfs3_3_0_base_u1 = index.add %rfs3_3_0_row_u1, %blk_k_u1 : index + %rfs3_3_0_idx_u1 = index.assume %rfs3_3_0_base_u1 [range(%rfs3_3_0_base_u1, 0, 5112)] : index + %rfs3_3_0_raw_u1 = vector.load %wl_flat[%rfs3_3_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_3_0_words_u1 = vector.bitcast %rfs3_3_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_3_0_u1 = vector.fragment %rfs3_3_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs3_3_1_row_u1 = index.mul %ln7, %k_wstride : index + %rfs3_3_1_base_u1 = index.add %rfs3_3_1_row_u1, %blk_k16_u1 : index + %rfs3_3_1_idx_u1 = index.assume %rfs3_3_1_base_u1 [range(%rfs3_3_1_base_u1, 0, 5112)] : index + %rfs3_3_1_raw_u1 = vector.load %wl_flat[%rfs3_3_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs3_3_1_words_u1 = vector.bitcast %rfs3_3_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs3_3_1_u1 = vector.fragment %rfs3_3_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_3_3_u1 = vector.mma %lf2_0_u1, %rfs3_3_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_3_3_u1 = vector.mma %lf2_1_u1, %rfs3_3_1_u1, %i0_3_3_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + + %ff3_3_u1 = vector.sitofp %i1_3_3_u1 : vector<8xi32> to vector<8xf32> + %sv3_3_u1 = vector.mulf %asv2_u0, %wsv7_u1 : vector<8xf32> + %fn3_3 = vector.fmaf %ff3_3_u1, %sv3_3_u1, %fn3_3_u0 : vector<8xf32> + scf.schedule.fence + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %fn0_0, %fn0_1, %fn0_2, %fn0_3, %fn1_0, %fn1_1, %fn1_2, %fn1_3, %fn2_0, %fn2_1, %fn2_2, %fn2_3, %fn3_0, %fn3_1, %fn3_2, %fn3_3 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + + %gate_m0_n0 = vector.siluf %f0_0 : vector<8xf32> + %gate_m0_n1 = vector.siluf %f0_1 : vector<8xf32> + %gate_m1_n0 = vector.siluf %f1_0 : vector<8xf32> + %gate_m1_n1 = vector.siluf %f1_1 : vector<8xf32> + %out_m0_n0 = vector.mulf %gate_m0_n0, %f0_2 : vector<8xf32> + %out_m0_n1 = vector.mulf %gate_m0_n1, %f0_3 : vector<8xf32> + %out_m1_n0 = vector.mulf %gate_m1_n0, %f1_2 : vector<8xf32> + %out_m1_n1 = vector.mulf %gate_m1_n1, %f1_3 : vector<8xf32> + %gate_m2_n0 = vector.siluf %f2_0 : vector<8xf32> + %gate_m2_n1 = vector.siluf %f2_1 : vector<8xf32> + %gate_m3_n0 = vector.siluf %f3_0 : vector<8xf32> + %gate_m3_n1 = vector.siluf %f3_1 : vector<8xf32> + %out_m2_n0 = vector.mulf %gate_m2_n0, %f2_2 : vector<8xf32> + %out_m2_n1 = vector.mulf %gate_m2_n1, %f2_3 : vector<8xf32> + %out_m3_n0 = vector.mulf %gate_m3_n0, %f3_2 : vector<8xf32> + %out_m3_n1 = vector.mulf %gate_m3_n1, %f3_3 : vector<8xf32> + + scf.if %publish_f32 { + vector.fragment.store %out_m0_n0, %dst_view[%gm0, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m0_n1, %dst_view[%gm0, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m1_n0, %dst_view[%gm0, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m1_n1, %dst_view[%gm0, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m2_n0, %dst_view[%gm2, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m2_n1, %dst_view[%gm2, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m3_n0, %dst_view[%gm2, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out_m3_n1, %dst_view[%gm2, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + } + + // Materialize the Q6_K down input through LDS for contiguous F16 stores and matching conversion order. + scf.if %publish_f16 { + %f16_tile = buffer.view %qout_scratch[%base] : buffer -> view<128x32xf32> + %f16_c4 = index.constant 4 : index + %f16_c8 = index.constant 8 : index + %f16_c32 = index.constant 32 : index + %f16_c512 = index.constant 512 : index + vector.fragment.store %out_m0_n0, %f16_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out_m0_n1, %f16_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out_m1_n0, %f16_tile[%lm1, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out_m1_n1, %f16_tile[%lm1, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %f16_batch = [%c0 to %f16_c8 step %c1] { + %f16_batch_base = index.mul %f16_batch, %f16_c512 : index + %f16_thread_base = index.mul %tid, %f16_c4 : index + %f16_linear0 = index.add %f16_batch_base, %f16_thread_base : index + %f16_linear = index.assume %f16_linear0 [range(%f16_linear0, 0, 4092), mul(%f16_linear0, 4)] : index + %f16_token_local = index.div %f16_linear, %f16_c32 : index + %f16_row_local = index.rem %f16_linear, %f16_c32 : index + %f16_wide = vector.load %f16_tile[%f16_token_local, %f16_row_local] : view<128x32xf32> -> vector<4xf32> + %f16_narrow = vector.fptrunc %f16_wide : vector<4xf32> to vector<4xf16> + %f16_token0 = index.add %col_base, %f16_token_local : index + %f16_row0 = index.add %row_base, %f16_row_local : index + %f16_row_end0 = index.add %f16_row0, %f16_c4 : index + %f16_token, %f16_row, %f16_row_end, %f16_cols, %f16_rows = index.assume %f16_token0, %f16_row0, %f16_row_end0, %cols_b, %rows_b [lt(%f16_token0, %cols_b), le(%f16_row_end0, %rows_b)] : index, index, index, index, index + vector.store %f16_narrow, %dst_f16_view[%f16_token, %f16_row] : vector<4xf16>, view<[%cols_b]x[%rows_b]xf16> + } + } + + // Publish affine-U4 K32 blocks while the SwiGLU tile is live to avoid an F32 round trip and separate quantizer. + scf.if %publish_u4 { + %u4_tile = buffer.view %qout_scratch[%base] : buffer -> view<128x32xf32> + + %u4_c4 = index.constant 4 : index + %u4_c8 = index.constant 8 : index + %u4_c16 = index.constant 16 : index + %u4_c32 = index.constant 32 : index + %u4_zero_f32 = scalar.constant 0.0 : f32 + %u4_zero_i32 = scalar.constant 0 : i32 + %u4_eps = scalar.constant 1.0000000000000001e-30 : f32 + %u4_f15 = scalar.constant 15.0 : f32 + %u4_vzero = vector.constant 0.0 : vector<4xf32> + %u4_v15 = vector.constant 15.0 : vector<4xf32> + %u4_bias_neg8 = vector.constant -8 : vector<4xi32> + %u4_nibble_mask = vector.constant 15 : vector<4xi32> + %u4_c8_i32 = scalar.constant 8 : i32 + %u4_xor1 = scalar.constant 1 : i32 + %u4_xor2 = scalar.constant 2 : i32 + %u4_xor4 = scalar.constant 4 : i32 + %u4_shift4 = scalar.constant 4 : i32 + %u4_shift8 = scalar.constant 8 : i32 + %u4_shift12 = scalar.constant 12 : i32 + %u4_shuffle_width = scalar.constant 32 : i32 + + %u4_total = index.mul %rows_b, %cols_b : index + %u4_groups0 = index.div %u4_total, %u4_c32 : index + %u4_groups = index.assume %u4_groups0 [range(%u4_groups0, 1, 33554432)] : index + %u4_halfwords = index.div %u4_total, %u4_c4 : index + %u4_qs = buffer.view %qout_qs_na[%base] : buffer -> view<[%u4_halfwords]xi16> + %u4_ds = buffer.view %qout_ds_na[%base] : buffer -> view<[%u4_groups]xf32> + %u4_meta_count = index.mul %u4_groups, %c2 : index + %u4_meta = buffer.view %qout_sums_na[%base] : buffer -> view<[%u4_meta_count]xi32> + + %u4_group_in_wave = index.div %lane, %u4_c8 : index + %u4_lane_in_group = index.rem %lane, %u4_c8 : index + %u4_is_leader = index.cmp eq, %u4_lane_in_group, %c0 : index + %u4_wave_group_base = index.mul %wave, %c4 : index + %u4_local_group_base = index.add %u4_wave_group_base, %u4_group_in_wave : index + %u4_row_group = index.assume %row_tile [range(%row_tile, 0, 8191)] : index + %u4_row_groups = index.div %rows_b, %u4_c32 : index + + vector.fragment.store %out_m0_n0, %u4_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out_m0_n1, %u4_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out_m1_n0, %u4_tile[%lm1, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out_m1_n1, %u4_tile[%lm1, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + + scf.for %u4_batch = [%c0 to %u4_c8 step %c1] { + %u4_batch_group = index.mul %u4_batch, %u4_c16 : index + %u4_local_group0 = index.add %u4_batch_group, %u4_local_group_base : index + %u4_token_local = index.assume %u4_local_group0 [range(%u4_local_group0, 0, 127)] : index + %u4_lane_row0 = index.mul %u4_lane_in_group, %u4_c4 : index + %u4_lane_row = index.assume %u4_lane_row0 [range(%u4_lane_row0, 0, 28)] : index + %u4_values = vector.load %u4_tile[%u4_token_local, %u4_lane_row] : view<128x32xf32> -> vector<4xf32> + + %u4_lane_max = vector.reduce %u4_values, %u4_zero_f32 : vector<4xf32>, f32 + %u4_neg_values = vector.subf %u4_vzero, %u4_values : vector<4xf32> + %u4_lane_neg_min = vector.reduce %u4_neg_values, %u4_zero_f32 : vector<4xf32>, f32 + %u4_max1_peer, %u4_max1_valid = kernel.subgroup.shuffle %u4_lane_max, %u4_xor1, %u4_shuffle_width : f32, i32, i32 + %u4_max1 = scalar.maxnumf %u4_lane_max, %u4_max1_peer : f32 + %u4_max2_peer, %u4_max2_valid = kernel.subgroup.shuffle %u4_max1, %u4_xor2, %u4_shuffle_width : f32, i32, i32 + %u4_max2 = scalar.maxnumf %u4_max1, %u4_max2_peer : f32 + %u4_max4_peer, %u4_max4_valid = kernel.subgroup.shuffle %u4_max2, %u4_xor4, %u4_shuffle_width : f32, i32, i32 + %u4_group_max = scalar.maxnumf %u4_max2, %u4_max4_peer : f32 + %u4_min1_peer, %u4_min1_valid = kernel.subgroup.shuffle %u4_lane_neg_min, %u4_xor1, %u4_shuffle_width : f32, i32, i32 + %u4_min1 = scalar.maxnumf %u4_lane_neg_min, %u4_min1_peer : f32 + %u4_min2_peer, %u4_min2_valid = kernel.subgroup.shuffle %u4_min1, %u4_xor2, %u4_shuffle_width : f32, i32, i32 + %u4_min2 = scalar.maxnumf %u4_min1, %u4_min2_peer : f32 + %u4_min4_peer, %u4_min4_valid = kernel.subgroup.shuffle %u4_min2, %u4_xor4, %u4_shuffle_width : f32, i32, i32 + %u4_group_neg_min = scalar.maxnumf %u4_min2, %u4_min4_peer : f32 + %u4_range0 = scalar.addf %u4_group_max, %u4_group_neg_min : f32 + %u4_range = scalar.maxnumf %u4_range0, %u4_eps : f32 + %u4_scale = scalar.divf %u4_range, %u4_f15 : f32 + %u4_rscale = scalar.divf %u4_f15, %u4_range : f32 + %u4_zp_raw = scalar.mulf %u4_group_neg_min, %u4_rscale : f32 + %u4_zp_round = scalar.roundf %u4_zp_raw : f32 + %u4_zp_clamped = scalar.clampf %u4_zp_round, %u4_zero_f32, %u4_f15 : f32 + %u4_zp = scalar.fptosi %u4_zp_clamped : f32 to i32 + + %u4_token = index.add %col_base, %u4_token_local : index + %u4_gid_base = index.mul %u4_token, %u4_row_groups : index + %u4_gid0 = index.add %u4_gid_base, %u4_row_group : index + %u4_gid = index.assume %u4_gid0 [range(%u4_gid0, 0, 33554431)] : index + scf.if %u4_is_leader { + view.store %u4_scale, %u4_ds[%u4_gid] : f32, view<[%u4_groups]xf32> + } + + %u4_rs = vector.splat %u4_rscale : vector<4xf32> + %u4_zpv = vector.splat %u4_zp_clamped : vector<4xf32> + %u4_scaled = vector.mulf %u4_values, %u4_rs : vector<4xf32> + %u4_shifted = vector.addf %u4_scaled, %u4_zpv : vector<4xf32> + %u4_rounded = vector.roundf %u4_shifted : vector<4xf32> + %u4_nonnegative = vector.maxnumf %u4_rounded, %u4_vzero : vector<4xf32> + %u4_to_ceiling = vector.subf %u4_v15, %u4_nonnegative : vector<4xf32> + %u4_ceiling_nonnegative = vector.maxnumf %u4_to_ceiling, %u4_vzero : vector<4xf32> + %u4_clamped = vector.subf %u4_v15, %u4_ceiling_nonnegative : vector<4xf32> + %u4_q32_unsigned = vector.fptosi %u4_clamped : vector<4xf32> to vector<4xi32> + %u4_q32 = vector.addi %u4_q32_unsigned, %u4_bias_neg8 : vector<4xi32> + %u4_qbits = vector.andi %u4_q32, %u4_nibble_mask : vector<4xi32> + %u4_q0 = vector.extract %u4_qbits[0] : vector<4xi32> -> i32 + %u4_q1 = vector.extract %u4_qbits[1] : vector<4xi32> -> i32 + %u4_q2 = vector.extract %u4_qbits[2] : vector<4xi32> -> i32 + %u4_q3 = vector.extract %u4_qbits[3] : vector<4xi32> -> i32 + %u4_q1s = scalar.shli %u4_q1, %u4_shift4 : i32 + %u4_q2s = scalar.shli %u4_q2, %u4_shift8 : i32 + %u4_q3s = scalar.shli %u4_q3, %u4_shift12 : i32 + %u4_q01 = scalar.ori %u4_q0, %u4_q1s : i32 + %u4_q23 = scalar.ori %u4_q2s, %u4_q3s : i32 + %u4_packed32 = scalar.ori %u4_q01, %u4_q23 : i32 + %u4_packed = scalar.trunci %u4_packed32 : i32 to i16 + %u4_word_base = index.mul %u4_gid, %u4_c8 : index + %u4_word_index0 = index.add %u4_word_base, %u4_lane_in_group : index + %u4_word_index = index.assume %u4_word_index0 [range(%u4_word_index0, 0, 268435455)] : index + view.store %u4_packed, %u4_qs[%u4_word_index] : i16, view<[%u4_halfwords]xi16> + + %u4_lane_sum = vector.reduce %u4_q32, %u4_zero_i32 : vector<4xi32>, i32 + %u4_sum1_peer, %u4_sum1_valid = kernel.subgroup.shuffle %u4_lane_sum, %u4_xor1, %u4_shuffle_width : i32, i32, i32 + %u4_sum1 = scalar.addi %u4_lane_sum, %u4_sum1_peer : i32 + %u4_sum2_peer, %u4_sum2_valid = kernel.subgroup.shuffle %u4_sum1, %u4_xor2, %u4_shuffle_width : i32, i32, i32 + %u4_sum2 = scalar.addi %u4_sum1, %u4_sum2_peer : i32 + %u4_sum4_peer, %u4_sum4_valid = kernel.subgroup.shuffle %u4_sum2, %u4_xor4, %u4_shuffle_width : i32, i32, i32 + %u4_group_sum = scalar.addi %u4_sum2, %u4_sum4_peer : i32 + scf.if %u4_is_leader { + %u4_zp_signed = scalar.subi %u4_zp, %u4_c8_i32 : i32 + %u4_meta0 = index.mul %u4_gid, %c2 : index + %u4_meta1 = index.add %u4_meta0, %c1 : index + view.store %u4_group_sum, %u4_meta[%u4_meta0] : i32, view<[%u4_meta_count]xi32> + view.store %u4_zp_signed, %u4_meta[%u4_meta1] : i32, view<[%u4_meta_count]xi32> + } + } + } + + // Publish the M128xN64 SwiGLU tile as two block-major Q8 K32 planes. + scf.if %publish_q8_plane { + %q8_tile = buffer.view %qout_scratch[%base] : buffer -> view<64x64xf32> + %q8_c4 = index.constant 4 : index + %q8_c8 = index.constant 8 : index + %q8_c16 = index.constant 16 : index + %q8_c32 = index.constant 32 : index + %q8_c64 = index.constant 64 : index + %q8_zero_f32 = scalar.constant 0.0 : f32 + %q8_one_f32 = scalar.constant 1.0 : f32 + %q8_f127 = scalar.constant 127.0 : f32 + %q8_xor1 = scalar.constant 1 : i32 + %q8_xor2 = scalar.constant 2 : i32 + %q8_xor4 = scalar.constant 4 : i32 + %q8_shuffle_width = scalar.constant 32 : i32 + + %q8_group_in_wave = index.div %lane, %q8_c8 : index + %q8_lane_in_block = index.rem %lane, %q8_c8 : index + %q8_is_leader = index.cmp eq, %q8_lane_in_block, %c0 : index + %q8_wave_group_base = index.mul %wave, %c4 : index + %q8_local_group_base = index.add %q8_wave_group_base, %q8_group_in_wave : index + %q8_lane_row_base = index.mul %q8_lane_in_block, %q8_c4 : index + %q8_payload_bytes = index.mul %rows_b, %cols_b : index + %q8_payload_words = index.div %q8_payload_bytes, %q8_c4 : index + %q8_payload_end = index.cast %q8_payload_bytes : index to offset + %q8_block_count = index.div %rows_b, %q8_c32 : index + %q8_metadata_count = index.mul %q8_block_count, %cols_b : index + %q8_payload = buffer.view %qout_qs_na[%base] : buffer -> view<[%q8_payload_words]xi32> + %q8_metadata = buffer.view %qout_qs_na[%q8_payload_end] : buffer -> view<[%q8_metadata_count]xi32> + + vector.fragment.store %out_m0_n0, %q8_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + vector.fragment.store %out_m0_n1, %q8_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + vector.fragment.store %out_m1_n0, %q8_tile[%lm0, %c32] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + %q8_c48 = index.constant 48 : index + vector.fragment.store %out_m1_n1, %q8_tile[%lm0, %q8_c48] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + + scf.for %q8_batch = [%c0 to %q8_c8 step %c1] { + %q8_batch_group = index.mul %q8_batch, %q8_c16 : index + %q8_local_group0 = index.add %q8_batch_group, %q8_local_group_base : index + %q8_local_group = index.assume %q8_local_group0 [range(%q8_local_group0, 0, 127)] : index + %q8_token_local0 = index.div %q8_local_group, %c2 : index + %q8_token_local = index.assume %q8_token_local0 [range(%q8_token_local0, 0, 63)] : index + %q8_row_block0 = index.rem %q8_local_group, %c2 : index + %q8_row_block = index.assume %q8_row_block0 [range(%q8_row_block0, 0, 1)] : index + %q8_row_block_base = index.mul %q8_row_block, %q8_c32 : index + %q8_lane_row0 = index.add %q8_row_block_base, %q8_lane_row_base : index + %q8_lane_row = index.assume %q8_lane_row0 [range(%q8_lane_row0, 0, 60), mul(%q8_lane_row0, 4)] : index + %q8_values = vector.load %q8_tile[%q8_token_local, %q8_lane_row] : view<64x64xf32> -> vector<4xf32> + %q8_absolute_values = vector.absf %q8_values : vector<4xf32> + %q8_lane_max = vector.reduce %q8_absolute_values, %q8_zero_f32 : vector<4xf32>, f32 + %q8_max1_peer, %q8_max1_valid = kernel.subgroup.shuffle %q8_lane_max, %q8_xor1, %q8_shuffle_width : f32, i32, i32 + %q8_max1 = scalar.maxnumf %q8_lane_max, %q8_max1_peer : f32 + %q8_max2_peer, %q8_max2_valid = kernel.subgroup.shuffle %q8_max1, %q8_xor2, %q8_shuffle_width : f32, i32, i32 + %q8_max2 = scalar.maxnumf %q8_max1, %q8_max2_peer : f32 + %q8_max4_peer, %q8_max4_valid = kernel.subgroup.shuffle %q8_max2, %q8_xor4, %q8_shuffle_width : f32, i32, i32 + %q8_amax = scalar.maxnumf %q8_max2, %q8_max4_peer : f32 + %q8_d = scalar.divf %q8_amax, %q8_f127 : f32 + %q8_d_nonzero = scalar.cmpf one, %q8_d, %q8_zero_f32 : f32 + %q8_d_inverse = scf.if %q8_d_nonzero -> (f32) { + %q8_inverse = scalar.divf %q8_one_f32, %q8_d : f32 + scf.yield %q8_inverse : f32 + } else { + scf.yield %q8_zero_f32 : f32 + } + %q8_d_inverse_vector = vector.splat %q8_d_inverse : vector<4xf32> + %q8_scaled_values = vector.mulf %q8_values, %q8_d_inverse_vector : vector<4xf32> + %q8_rounded_values = vector.roundf %q8_scaled_values : vector<4xf32> + %q8_quantized_values = vector.fptosi %q8_rounded_values : vector<4xf32> to vector<4xi8> + %q8_packed_word = vector.bitcast %q8_quantized_values : vector<4xi8> to vector<1xi32> + + %q8_token = index.add %col_base, %q8_token_local : index + %q8_global_block_base = index.mul %row_tile, %c2 : index + %q8_global_block0 = index.add %q8_global_block_base, %q8_row_block : index + %q8_global_block = index.assume %q8_global_block0 [range(%q8_global_block0, 0, 16383)] : index + %q8_block_token_base = index.mul %q8_global_block, %cols_b : index + %q8_block_token0 = index.add %q8_block_token_base, %q8_token : index + %q8_block_token = index.assume %q8_block_token0 [range(%q8_block_token0, 0, 268435455)] : index + %q8_word_base = index.mul %q8_block_token, %q8_c8 : index + %q8_word_index0 = index.add %q8_word_base, %q8_lane_in_block : index + %q8_word_index = index.assume %q8_word_index0 [range(%q8_word_index0, 0, 2147483647)] : index + vector.store %q8_packed_word, %q8_payload[%q8_word_index] : vector<1xi32>, view<[%q8_payload_words]xi32> + + %q8_lane_sum = vector.reduce %q8_rounded_values, %q8_zero_f32 : vector<4xf32>, f32 + %q8_sum1_peer, %q8_sum1_valid = kernel.subgroup.shuffle %q8_lane_sum, %q8_xor1, %q8_shuffle_width : f32, i32, i32 + %q8_sum1 = scalar.addf %q8_lane_sum, %q8_sum1_peer : f32 + %q8_sum2_peer, %q8_sum2_valid = kernel.subgroup.shuffle %q8_sum1, %q8_xor2, %q8_shuffle_width : f32, i32, i32 + %q8_sum2 = scalar.addf %q8_sum1, %q8_sum2_peer : f32 + %q8_sum4_peer, %q8_sum4_valid = kernel.subgroup.shuffle %q8_sum2, %q8_xor4, %q8_shuffle_width : f32, i32, i32 + %q8_quantized_sum = scalar.addf %q8_sum2, %q8_sum4_peer : f32 + scf.if %q8_is_leader { + %q8_s = scalar.mulf %q8_quantized_sum, %q8_d : f32 + %q8_d_f16 = scalar.fptrunc %q8_d : f32 to f16 + %q8_s_f16 = scalar.fptrunc %q8_s : f32 to f16 + %q8_ds_pair = vector.from_elements %q8_d_f16, %q8_s_f16 : vector<2xf16> + %q8_ds_word = vector.bitcast %q8_ds_pair : vector<2xf16> to vector<1xi32> + vector.store %q8_ds_word, %q8_metadata[%q8_block_token] : vector<1xi32>, view<[%q8_metadata_count]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + vector.fragment.store %out_m2_n0, %q8_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + vector.fragment.store %out_m2_n1, %q8_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + vector.fragment.store %out_m3_n0, %q8_tile[%lm0, %q8_c32] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + vector.fragment.store %out_m3_n1, %q8_tile[%lm0, %q8_c48] shape [%c16, %c16] : vector<8xf32>, view<64x64xf32> + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %q8b_batch = [%c0 to %q8_c8 step %c1] { + %q8b_batch_group = index.mul %q8b_batch, %q8_c16 : index + %q8b_local_group0 = index.add %q8b_batch_group, %q8_local_group_base : index + %q8b_local_group = index.assume %q8b_local_group0 [range(%q8b_local_group0, 0, 127)] : index + %q8b_token_local0 = index.div %q8b_local_group, %c2 : index + %q8b_token_local = index.assume %q8b_token_local0 [range(%q8b_token_local0, 0, 63)] : index + %q8b_row_block0 = index.rem %q8b_local_group, %c2 : index + %q8b_row_block = index.assume %q8b_row_block0 [range(%q8b_row_block0, 0, 1)] : index + %q8b_row_block_base = index.mul %q8b_row_block, %q8_c32 : index + %q8b_lane_row0 = index.add %q8b_row_block_base, %q8_lane_row_base : index + %q8b_lane_row = index.assume %q8b_lane_row0 [range(%q8b_lane_row0, 0, 60), mul(%q8b_lane_row0, 4)] : index + %q8b_values = vector.load %q8_tile[%q8b_token_local, %q8b_lane_row] : view<64x64xf32> -> vector<4xf32> + %q8b_absolute_values = vector.absf %q8b_values : vector<4xf32> + %q8b_lane_max = vector.reduce %q8b_absolute_values, %q8_zero_f32 : vector<4xf32>, f32 + %q8b_max1_peer, %q8b_max1_valid = kernel.subgroup.shuffle %q8b_lane_max, %q8_xor1, %q8_shuffle_width : f32, i32, i32 + %q8b_max1 = scalar.maxnumf %q8b_lane_max, %q8b_max1_peer : f32 + %q8b_max2_peer, %q8b_max2_valid = kernel.subgroup.shuffle %q8b_max1, %q8_xor2, %q8_shuffle_width : f32, i32, i32 + %q8b_max2 = scalar.maxnumf %q8b_max1, %q8b_max2_peer : f32 + %q8b_max4_peer, %q8b_max4_valid = kernel.subgroup.shuffle %q8b_max2, %q8_xor4, %q8_shuffle_width : f32, i32, i32 + %q8b_amax = scalar.maxnumf %q8b_max2, %q8b_max4_peer : f32 + %q8b_d = scalar.divf %q8b_amax, %q8_f127 : f32 + %q8b_d_nonzero = scalar.cmpf one, %q8b_d, %q8_zero_f32 : f32 + %q8b_d_inverse = scf.if %q8b_d_nonzero -> (f32) { + %q8b_inverse = scalar.divf %q8_one_f32, %q8b_d : f32 + scf.yield %q8b_inverse : f32 + } else { + scf.yield %q8_zero_f32 : f32 + } + %q8b_d_inverse_vector = vector.splat %q8b_d_inverse : vector<4xf32> + %q8b_scaled_values = vector.mulf %q8b_values, %q8b_d_inverse_vector : vector<4xf32> + %q8b_rounded_values = vector.roundf %q8b_scaled_values : vector<4xf32> + %q8b_quantized_values = vector.fptosi %q8b_rounded_values : vector<4xf32> to vector<4xi8> + %q8b_packed_word = vector.bitcast %q8b_quantized_values : vector<4xi8> to vector<1xi32> + + %q8b_token_within_tile = index.add %q8b_token_local, %q8_c64 : index + %q8b_token = index.add %col_base, %q8b_token_within_tile : index + %q8b_global_block_base = index.mul %row_tile, %c2 : index + %q8b_global_block0 = index.add %q8b_global_block_base, %q8b_row_block : index + %q8b_global_block = index.assume %q8b_global_block0 [range(%q8b_global_block0, 0, 16383)] : index + %q8b_block_token_base = index.mul %q8b_global_block, %cols_b : index + %q8b_block_token0 = index.add %q8b_block_token_base, %q8b_token : index + %q8b_block_token = index.assume %q8b_block_token0 [range(%q8b_block_token0, 0, 268435455)] : index + %q8b_word_base = index.mul %q8b_block_token, %q8_c8 : index + %q8b_word_index0 = index.add %q8b_word_base, %q8_lane_in_block : index + %q8b_word_index = index.assume %q8b_word_index0 [range(%q8b_word_index0, 0, 2147483647)] : index + vector.store %q8b_packed_word, %q8_payload[%q8b_word_index] : vector<1xi32>, view<[%q8_payload_words]xi32> + + %q8b_lane_sum = vector.reduce %q8b_rounded_values, %q8_zero_f32 : vector<4xf32>, f32 + %q8b_sum1_peer, %q8b_sum1_valid = kernel.subgroup.shuffle %q8b_lane_sum, %q8_xor1, %q8_shuffle_width : f32, i32, i32 + %q8b_sum1 = scalar.addf %q8b_lane_sum, %q8b_sum1_peer : f32 + %q8b_sum2_peer, %q8b_sum2_valid = kernel.subgroup.shuffle %q8b_sum1, %q8_xor2, %q8_shuffle_width : f32, i32, i32 + %q8b_sum2 = scalar.addf %q8b_sum1, %q8b_sum2_peer : f32 + %q8b_sum4_peer, %q8b_sum4_valid = kernel.subgroup.shuffle %q8b_sum2, %q8_xor4, %q8_shuffle_width : f32, i32, i32 + %q8b_quantized_sum = scalar.addf %q8b_sum2, %q8b_sum4_peer : f32 + scf.if %q8_is_leader { + %q8b_s = scalar.mulf %q8b_quantized_sum, %q8b_d : f32 + %q8b_d_f16 = scalar.fptrunc %q8b_d : f32 to f16 + %q8b_s_f16 = scalar.fptrunc %q8b_s : f32 to f16 + %q8b_ds_pair = vector.from_elements %q8b_d_f16, %q8b_s_f16 : vector<2xf16> + %q8b_ds_word = vector.bitcast %q8b_ds_pair : vector<2xf16> to vector<1xi32> + vector.store %q8b_ds_word, %q8_metadata[%q8b_block_token] : vector<1xi32>, view<[%q8_metadata_count]xi32> + } + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_symmetric_i4_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_symmetric_i4_wmma.loom new file mode 100644 index 000000000000..af363bf16e04 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_symmetric_i4_wmma.loom @@ -0,0 +1,3451 @@ +template.decl @ggml.quantize_symmetric_i4_k64.body(%plane_major: i1, %src: buffer, %qs: buffer, %ds: buffer, %sums: buffer) + +config.decl @ggml.quantize_symmetric_i4_k64.input_size : %value: index where [range(%value, 32, 65536), mul(%value, 32)] + +config.decl @ggml.quantize_symmetric_i4_k64.token_count : %value: index where [range(%value, 1, 16384)] + +kernel.def export("ggml_quantize_f32_symmetric_i4_k64_plane") @ggml_quantize_f32_symmetric_i4_k64_plane() { + %unit = index.constant 1 : index + %wg = index.constant 256 : index + %wg_groups_m1 = index.constant 127 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %k = config.get @ggml.quantize_symmetric_i4_k64.input_size : index + %cols = config.get @ggml.quantize_symmetric_i4_k64.token_count : index + %total = index.mul %k, %cols : index + %groups = index.div %total, %c64 : index + %rounded = index.add %groups, %wg_groups_m1 : index + %workgroups = index.div %rounded, %c128 : index + kernel.launch.config workgroups(%workgroups, %unit, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%src: buffer, %qs: buffer, %ds: buffer, %sums: buffer) { + %plane_major = scalar.constant true : i1 + template.apply<@ggml.quantize_symmetric_i4_k64.body>(%plane_major, %src, %qs, %ds, %sums) : (i1, buffer, buffer, buffer, buffer) + kernel.return +} + +// Exact low-column direct-dot specializations for adjacent projections. +func.def inline @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words: vector<2xi32>, %group: index, %token: index, %half: index, %k: index, %group_count: index, %qact: buffer, %scales: buffer) -> (i32, f32) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %zero_i32 = scalar.constant 0 : i32 + %zero_i32x2 = vector.constant 0 : vector<2xi32> + %qact_flat = buffer.view %qact[%base] : buffer -> view<1073741824xi8> + %scale_flat = buffer.view %scales[%base] : buffer -> view<33554432xf32> + + %token_bytes = index.div %k, %c2 : index + %token_payload_base = index.mul %token, %token_bytes : index + %group_payload_add = index.mul %group, %c16 : index + %half_payload_add = index.mul %half, %c8 : index + %payload_index0 = index.add %token_payload_base, %group_payload_add : index + %payload_index1 = index.add %payload_index0, %half_payload_add : index + %payload_index = index.assume %payload_index1 [range(%payload_index1, 0, 1073741816)] : index + %activation_bytes = vector.load %qact_flat[%payload_index] : view<1073741824xi8> -> vector<8xi8> + %activation_words = vector.bitcast %activation_bytes : vector<8xi8> to vector<2xi32> + %dot_parts = vector.dot8i4 %weight_words, %activation_words, %zero_i32x2 : vector<2xi32> + %dot = vector.reduce %dot_parts, %zero_i32 : vector<2xi32>, i32 + + %token_scale_base = index.mul %token, %group_count : index + %scale_index0 = index.add %token_scale_base, %group : index + %scale_index = index.assume %scale_index0 [range(%scale_index0, 0, 33554431)] : index + %activation_scale = view.load %scale_flat[%scale_index] : view<33554432xf32> -> f32 + func.return %dot, %activation_scale : i32, f32 +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c1") @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c1() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %rows_up = index.add %rows, %c15 : index + %row_workgroups = index.div %rows_up, %c16 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%gate_weight: buffer, %up_weight: buffer, %gate_output: buffer, %up_output: buffer, %qact: buffer, %scales: buffer, %sums: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 1, 1)] : index + %base = index.constant 0 : offset + %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 + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + %xor_one = scalar.constant 1 : i32 + %shuffle_width = scalar.constant 32 : i32 + + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %projection = index.rem %workgroup0, %c2 : index + %workgroup = index.div %workgroup0, %c2 : index + %is_up = index.cmp eq, %projection, %c1 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 31)] : index + %row_in_wave = index.div %workitem, %c2 : index + %half = index.rem %workitem, %c2 : index + %row_base = index.mul %workgroup, %c16 : index + %row = index.add %row_base, %row_in_wave : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %gate_weight_global = buffer.assume.memory_space %gate_weight : buffer + %up_weight_global = buffer.assume.memory_space %up_weight : buffer + %gate_output_global = buffer.assume.memory_space %gate_output : buffer + %up_output_global = buffer.assume.memory_space %up_output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %weight_global = scf.select %is_up, %up_weight_global, %gate_weight_global : buffer + %output_global = scf.select %is_up, %up_output_global, %gate_output_global : buffer + %weight_na, %output_na, %qact_na, %scales_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global : buffer, buffer, buffer, buffer + + %result0 = scf.if %valid_row -> (f32) { + %acc0 = scf.for %group = [%c0 to %group_count step %c1](%iter0 = %zero_f32 : f32) -> (f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %half_payload_add = index.mul %half, %c8 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index2 = index.add %payload_index1, %row_payload_add : index + %payload_index = index.add %payload_index2, %half_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<2xi32> + %weight_words = vector.load %weight_view[%c0] : view<2xi32> -> vector<2xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %dot0, %as0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c0, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %peer0, %valid0 = kernel.subgroup.shuffle %dot0, %xor_one, %shuffle_width : i32, i32, i32 + %full0 = scalar.addi %dot0, %peer0 : i32 + %fp0 = scalar.sitofp %full0 : i32 to f32 + %scale0 = scalar.mulf %weight_scale, %as0 : f32 + %next0 = scalar.fmaf %fp0, %scale0, %iter0 : f32 + scf.yield %next0 : f32 + } + scf.yield %acc0 : f32 + } else { + scf.yield %zero_f32 : f32 + } + + %is_writer = index.cmp eq, %half, %c0 : index + %writes = scalar.andi %valid_row, %is_writer : i1 + scf.if %writes { + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c2") @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c2() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %rows_up = index.add %rows, %c15 : index + %row_workgroups = index.div %rows_up, %c16 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%gate_weight: buffer, %up_weight: buffer, %gate_output: buffer, %up_output: buffer, %qact: buffer, %scales: buffer, %sums: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 2, 2)] : index + %base = index.constant 0 : offset + %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 + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + %xor_one = scalar.constant 1 : i32 + %shuffle_width = scalar.constant 32 : i32 + + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %projection = index.rem %workgroup0, %c2 : index + %workgroup = index.div %workgroup0, %c2 : index + %is_up = index.cmp eq, %projection, %c1 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 31)] : index + %row_in_wave = index.div %workitem, %c2 : index + %half = index.rem %workitem, %c2 : index + %row_base = index.mul %workgroup, %c16 : index + %row = index.add %row_base, %row_in_wave : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %gate_weight_global = buffer.assume.memory_space %gate_weight : buffer + %up_weight_global = buffer.assume.memory_space %up_weight : buffer + %gate_output_global = buffer.assume.memory_space %gate_output : buffer + %up_output_global = buffer.assume.memory_space %up_output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %weight_global = scf.select %is_up, %up_weight_global, %gate_weight_global : buffer + %output_global = scf.select %is_up, %up_output_global, %gate_output_global : buffer + %weight_na, %output_na, %qact_na, %scales_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global : buffer, buffer, buffer, buffer + + %result0, %result1 = scf.if %valid_row -> (f32, f32) { + %acc0, %acc1 = scf.for %group = [%c0 to %group_count step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32) -> (f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %half_payload_add = index.mul %half, %c8 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index2 = index.add %payload_index1, %row_payload_add : index + %payload_index = index.add %payload_index2, %half_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<2xi32> + %weight_words = vector.load %weight_view[%c0] : view<2xi32> -> vector<2xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %dot0, %as0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c0, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot1, %as1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c1, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %peer0, %valid0 = kernel.subgroup.shuffle %dot0, %xor_one, %shuffle_width : i32, i32, i32 + %peer1, %valid1 = kernel.subgroup.shuffle %dot1, %xor_one, %shuffle_width : i32, i32, i32 + %full0 = scalar.addi %dot0, %peer0 : i32 + %full1 = scalar.addi %dot1, %peer1 : i32 + %fp0 = scalar.sitofp %full0 : i32 to f32 + %fp1 = scalar.sitofp %full1 : i32 to f32 + %scale0 = scalar.mulf %weight_scale, %as0 : f32 + %scale1 = scalar.mulf %weight_scale, %as1 : f32 + %next0 = scalar.fmaf %fp0, %scale0, %iter0 : f32 + %next1 = scalar.fmaf %fp1, %scale1, %iter1 : f32 + scf.yield %next0, %next1 : f32, f32 + } + scf.yield %acc0, %acc1 : f32, f32 + } else { + scf.yield %zero_f32, %zero_f32 : f32, f32 + } + + %is_writer = index.cmp eq, %half, %c0 : index + %writes = scalar.andi %valid_row, %is_writer : i1 + scf.if %writes { + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c3") @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c3() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %rows_up = index.add %rows, %c15 : index + %row_workgroups = index.div %rows_up, %c16 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%gate_weight: buffer, %up_weight: buffer, %gate_output: buffer, %up_output: buffer, %qact: buffer, %scales: buffer, %sums: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 3, 3)] : index + %base = index.constant 0 : offset + %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 + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + %xor_one = scalar.constant 1 : i32 + %shuffle_width = scalar.constant 32 : i32 + + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %projection = index.rem %workgroup0, %c2 : index + %workgroup = index.div %workgroup0, %c2 : index + %is_up = index.cmp eq, %projection, %c1 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 31)] : index + %row_in_wave = index.div %workitem, %c2 : index + %half = index.rem %workitem, %c2 : index + %row_base = index.mul %workgroup, %c16 : index + %row = index.add %row_base, %row_in_wave : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %gate_weight_global = buffer.assume.memory_space %gate_weight : buffer + %up_weight_global = buffer.assume.memory_space %up_weight : buffer + %gate_output_global = buffer.assume.memory_space %gate_output : buffer + %up_output_global = buffer.assume.memory_space %up_output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %weight_global = scf.select %is_up, %up_weight_global, %gate_weight_global : buffer + %output_global = scf.select %is_up, %up_output_global, %gate_output_global : buffer + %weight_na, %output_na, %qact_na, %scales_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global : buffer, buffer, buffer, buffer + + %result0, %result1, %result2 = scf.if %valid_row -> (f32, f32, f32) { + %acc0, %acc1, %acc2 = scf.for %group = [%c0 to %group_count step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32) -> (f32, f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %half_payload_add = index.mul %half, %c8 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index2 = index.add %payload_index1, %row_payload_add : index + %payload_index = index.add %payload_index2, %half_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<2xi32> + %weight_words = vector.load %weight_view[%c0] : view<2xi32> -> vector<2xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %dot0, %as0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c0, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot1, %as1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c1, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot2, %as2 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c2, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %peer0, %valid0 = kernel.subgroup.shuffle %dot0, %xor_one, %shuffle_width : i32, i32, i32 + %peer1, %valid1 = kernel.subgroup.shuffle %dot1, %xor_one, %shuffle_width : i32, i32, i32 + %peer2, %valid2 = kernel.subgroup.shuffle %dot2, %xor_one, %shuffle_width : i32, i32, i32 + %full0 = scalar.addi %dot0, %peer0 : i32 + %full1 = scalar.addi %dot1, %peer1 : i32 + %full2 = scalar.addi %dot2, %peer2 : i32 + %fp0 = scalar.sitofp %full0 : i32 to f32 + %fp1 = scalar.sitofp %full1 : i32 to f32 + %fp2 = scalar.sitofp %full2 : i32 to f32 + %scale0 = scalar.mulf %weight_scale, %as0 : f32 + %scale1 = scalar.mulf %weight_scale, %as1 : f32 + %scale2 = scalar.mulf %weight_scale, %as2 : f32 + %next0 = scalar.fmaf %fp0, %scale0, %iter0 : f32 + %next1 = scalar.fmaf %fp1, %scale1, %iter1 : f32 + %next2 = scalar.fmaf %fp2, %scale2, %iter2 : f32 + scf.yield %next0, %next1, %next2 : f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2 : f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32 + } + + %is_writer = index.cmp eq, %half, %c0 : index + %writes = scalar.andi %valid_row, %is_writer : i1 + scf.if %writes { + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c4") @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c4() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %rows_up = index.add %rows, %c15 : index + %row_workgroups = index.div %rows_up, %c16 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%gate_weight: buffer, %up_weight: buffer, %gate_output: buffer, %up_output: buffer, %qact: buffer, %scales: buffer, %sums: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 4, 4)] : index + %base = index.constant 0 : offset + %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 + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + %xor_one = scalar.constant 1 : i32 + %shuffle_width = scalar.constant 32 : i32 + + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %projection = index.rem %workgroup0, %c2 : index + %workgroup = index.div %workgroup0, %c2 : index + %is_up = index.cmp eq, %projection, %c1 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 31)] : index + %row_in_wave = index.div %workitem, %c2 : index + %half = index.rem %workitem, %c2 : index + %row_base = index.mul %workgroup, %c16 : index + %row = index.add %row_base, %row_in_wave : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %gate_weight_global = buffer.assume.memory_space %gate_weight : buffer + %up_weight_global = buffer.assume.memory_space %up_weight : buffer + %gate_output_global = buffer.assume.memory_space %gate_output : buffer + %up_output_global = buffer.assume.memory_space %up_output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %weight_global = scf.select %is_up, %up_weight_global, %gate_weight_global : buffer + %output_global = scf.select %is_up, %up_output_global, %gate_output_global : buffer + %weight_na, %output_na, %qact_na, %scales_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global : buffer, buffer, buffer, buffer + + %result0, %result1, %result2, %result3 = scf.if %valid_row -> (f32, f32, f32, f32) { + %acc0, %acc1, %acc2, %acc3 = scf.for %group = [%c0 to %group_count step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32, %iter3 = %zero_f32 : f32) -> (f32, f32, f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %half_payload_add = index.mul %half, %c8 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index2 = index.add %payload_index1, %row_payload_add : index + %payload_index = index.add %payload_index2, %half_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<2xi32> + %weight_words = vector.load %weight_view[%c0] : view<2xi32> -> vector<2xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %dot0, %as0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c0, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot1, %as1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c1, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot2, %as2 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c2, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot3, %as3 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c3, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %peer0, %valid0 = kernel.subgroup.shuffle %dot0, %xor_one, %shuffle_width : i32, i32, i32 + %peer1, %valid1 = kernel.subgroup.shuffle %dot1, %xor_one, %shuffle_width : i32, i32, i32 + %peer2, %valid2 = kernel.subgroup.shuffle %dot2, %xor_one, %shuffle_width : i32, i32, i32 + %peer3, %valid3 = kernel.subgroup.shuffle %dot3, %xor_one, %shuffle_width : i32, i32, i32 + %full0 = scalar.addi %dot0, %peer0 : i32 + %full1 = scalar.addi %dot1, %peer1 : i32 + %full2 = scalar.addi %dot2, %peer2 : i32 + %full3 = scalar.addi %dot3, %peer3 : i32 + %fp0 = scalar.sitofp %full0 : i32 to f32 + %fp1 = scalar.sitofp %full1 : i32 to f32 + %fp2 = scalar.sitofp %full2 : i32 to f32 + %fp3 = scalar.sitofp %full3 : i32 to f32 + %scale0 = scalar.mulf %weight_scale, %as0 : f32 + %scale1 = scalar.mulf %weight_scale, %as1 : f32 + %scale2 = scalar.mulf %weight_scale, %as2 : f32 + %scale3 = scalar.mulf %weight_scale, %as3 : f32 + %next0 = scalar.fmaf %fp0, %scale0, %iter0 : f32 + %next1 = scalar.fmaf %fp1, %scale1, %iter1 : f32 + %next2 = scalar.fmaf %fp2, %scale2, %iter2 : f32 + %next3 = scalar.fmaf %fp3, %scale3, %iter3 : f32 + scf.yield %next0, %next1, %next2, %next3 : f32, f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2, %acc3 : f32, f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32, f32 + } + + %is_writer = index.cmp eq, %half, %c0 : index + %writes = scalar.andi %valid_row, %is_writer : i1 + scf.if %writes { + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result3, %output_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c5") @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_c5() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c15 = index.constant 15 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %rows_up = index.add %rows, %c15 : index + %row_workgroups = index.div %rows_up, %c16 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%gate_weight: buffer, %up_weight: buffer, %gate_output: buffer, %up_output: buffer, %qact: buffer, %scales: buffer, %sums: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 5, 5)] : index + %base = index.constant 0 : offset + %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 + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + %xor_one = scalar.constant 1 : i32 + %shuffle_width = scalar.constant 32 : i32 + + %workgroup0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %projection = index.rem %workgroup0, %c2 : index + %workgroup = index.div %workgroup0, %c2 : index + %is_up = index.cmp eq, %projection, %c1 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 31)] : index + %row_in_wave = index.div %workitem, %c2 : index + %half = index.rem %workitem, %c2 : index + %row_base = index.mul %workgroup, %c16 : index + %row = index.add %row_base, %row_in_wave : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %gate_weight_global = buffer.assume.memory_space %gate_weight : buffer + %up_weight_global = buffer.assume.memory_space %up_weight : buffer + %gate_output_global = buffer.assume.memory_space %gate_output : buffer + %up_output_global = buffer.assume.memory_space %up_output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %weight_global = scf.select %is_up, %up_weight_global, %gate_weight_global : buffer + %output_global = scf.select %is_up, %up_output_global, %gate_output_global : buffer + %weight_na, %output_na, %qact_na, %scales_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global : buffer, buffer, buffer, buffer + + %result0, %result1, %result2, %result3, %result4 = scf.if %valid_row -> (f32, f32, f32, f32, f32) { + %acc0, %acc1, %acc2, %acc3, %acc4 = scf.for %group = [%c0 to %group_count step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32, %iter3 = %zero_f32 : f32, %iter4 = %zero_f32 : f32) -> (f32, f32, f32, f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %half_payload_add = index.mul %half, %c8 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index2 = index.add %payload_index1, %row_payload_add : index + %payload_index = index.add %payload_index2, %half_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<2xi32> + %weight_words = vector.load %weight_view[%c0] : view<2xi32> -> vector<2xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %dot0, %as0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c0, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot1, %as1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c1, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot2, %as2 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c2, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot3, %as3 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c3, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %dot4, %as4 = func.call @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_direct_dot_half_dot(%weight_words, %group, %c4, %half, %k, %group_count, %qact_na, %scales_na) : (vector<2xi32>, index, index, index, index, index, buffer, buffer) -> (i32, f32) + %peer0, %valid0 = kernel.subgroup.shuffle %dot0, %xor_one, %shuffle_width : i32, i32, i32 + %peer1, %valid1 = kernel.subgroup.shuffle %dot1, %xor_one, %shuffle_width : i32, i32, i32 + %peer2, %valid2 = kernel.subgroup.shuffle %dot2, %xor_one, %shuffle_width : i32, i32, i32 + %peer3, %valid3 = kernel.subgroup.shuffle %dot3, %xor_one, %shuffle_width : i32, i32, i32 + %peer4, %valid4 = kernel.subgroup.shuffle %dot4, %xor_one, %shuffle_width : i32, i32, i32 + %full0 = scalar.addi %dot0, %peer0 : i32 + %full1 = scalar.addi %dot1, %peer1 : i32 + %full2 = scalar.addi %dot2, %peer2 : i32 + %full3 = scalar.addi %dot3, %peer3 : i32 + %full4 = scalar.addi %dot4, %peer4 : i32 + %fp0 = scalar.sitofp %full0 : i32 to f32 + %fp1 = scalar.sitofp %full1 : i32 to f32 + %fp2 = scalar.sitofp %full2 : i32 to f32 + %fp3 = scalar.sitofp %full3 : i32 to f32 + %fp4 = scalar.sitofp %full4 : i32 to f32 + %scale0 = scalar.mulf %weight_scale, %as0 : f32 + %scale1 = scalar.mulf %weight_scale, %as1 : f32 + %scale2 = scalar.mulf %weight_scale, %as2 : f32 + %scale3 = scalar.mulf %weight_scale, %as3 : f32 + %scale4 = scalar.mulf %weight_scale, %as4 : f32 + %next0 = scalar.fmaf %fp0, %scale0, %iter0 : f32 + %next1 = scalar.fmaf %fp1, %scale1, %iter1 : f32 + %next2 = scalar.fmaf %fp2, %scale2, %iter2 : f32 + %next3 = scalar.fmaf %fp3, %scale3, %iter3 : f32 + %next4 = scalar.fmaf %fp4, %scale4, %iter4 : f32 + scf.yield %next0, %next1, %next2, %next3, %next4 : f32, f32, f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2, %acc3, %acc4 : f32, f32, f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32, f32, f32 + } + + %is_writer = index.cmp eq, %half, %c0 : index + %writes = scalar.andi %valid_row, %is_writer : i1 + scf.if %writes { + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result3, %output_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result4, %output_view[%c4, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.return +} + +template.def<@ggml.quantize_symmetric_i4_k64.body> device @ggml_quantize_symmetric_i4_k64_body(%plane_major: i1, %src: buffer, %qs: buffer, %ds: buffer, %sums: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c64 = index.constant 64 : index + %c128 = index.constant 128 : index + %c0_f32 = scalar.constant 0.0 : f32 + %amax_eps = scalar.constant 1.0000000000000001e-30 : f32 + %inv7 = scalar.constant 0.14285714285714285 : f32 + %f7 = scalar.constant 7.0 : f32 + %xor1 = scalar.constant 1 : i32 + %xor2 = scalar.constant 2 : i32 + %xor4 = scalar.constant 4 : i32 + %xor8 = scalar.constant 8 : i32 + %shift4 = scalar.constant 4 : i32 + %shift8 = scalar.constant 8 : i32 + %shift12 = scalar.constant 12 : i32 + %shuffle_width = scalar.constant 32 : i32 + %mask15 = vector.constant 15 : vector<4xi32> + + %k = config.get @ggml.quantize_symmetric_i4_k64.input_size : index + %cols = config.get @ggml.quantize_symmetric_i4_k64.token_count : index + %k_b = index.assume %k [range(%k, 32, 65536)] : index + %cols_b = index.assume %cols [range(%cols, 1, 16384)] : index + %total = index.mul %k_b, %cols_b : index + %groups0 = index.div %total, %c32 : index + %groups = index.assume %groups0 [range(%groups0, 1, 33554432)] : index + %groups64_0 = index.div %total, %c64 : index + %groups64 = index.assume %groups64_0 [range(%groups64_0, 1, 16777216)] : index + %nchunks = index.div %k_b, %c64 : index + + %wgid0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %wgid = index.assume %wgid0 [range(%wgid0, 0, 131071)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %wave = index.div %tid, %c32 : index + %lane = index.rem %tid, %c32 : index + %group_in_wave = index.div %lane, %c16 : index + %lane_in_group = index.rem %lane, %c16 : index + %group_half = index.div %lane_in_group, %c8 : index + %lane_in_half = index.rem %lane_in_group, %c8 : index + %is_scale_leader = index.cmp eq, %lane_in_group, %c0 : index + %is_sum_leader = index.cmp eq, %lane_in_half, %c0 : index + %wg_group_base = index.mul %wgid, %c128 : index + %wave_group_base = index.mul %wave, %c2 : index + %local_group_base = index.add %wg_group_base, %wave_group_base : index + + %src_g = buffer.assume.memory_space %src : buffer + %qs_g = buffer.assume.memory_space %qs : buffer + %ds_g = buffer.assume.memory_space %ds : buffer + %sums_g = buffer.assume.memory_space %sums : buffer + %src_na, %qs_na, %ds_na, %sums_na = buffer.assume.noalias %src_g, %qs_g, %ds_g, %sums_g : buffer, buffer, buffer, buffer + %src_elems = index.assume %total [range(%total, 32, 1073741824)] : index + %src_view = buffer.view %src_na[%base] : buffer -> view<[%src_elems]xf32> + %qs_halfwords = index.div %total, %c4 : index + %qs_view = buffer.view %qs_na[%base] : buffer -> view<[%qs_halfwords]xi16> + %ds_view = buffer.view %ds_na[%base] : buffer -> view<[%groups]xf32> + %sum_view = buffer.view %sums_na[%base] : buffer -> view<[%groups]xi32> + + scf.for %batch = [%c0 to %c8 step %c1] { + %batch_group_off = index.mul %batch, %c16 : index + %gid0 = index.add %local_group_base, %batch_group_off : index + %gid1 = index.add %gid0, %group_in_wave : index + %gid = index.assume %gid1 [range(%gid1, 0, 16777215)] : index + %in_range = index.cmp ult, %gid, %groups64 : index + scf.if %in_range { + %ebase0 = index.mul %gid, %c64 : index + %lane_elem_off = index.mul %lane_in_group, %c4 : index + %ebase1 = index.add %ebase0, %lane_elem_off : index + %ebase = index.assume %ebase1 [range(%ebase1, 0, 1073741820)] : index + %v = vector.load %src_view[%ebase] : view<[%src_elems]xf32> -> vector<4xf32> + %av = vector.absf %v : vector<4xf32> + %lane_max = vector.reduce %av, %c0_f32 : vector<4xf32>, f32 + %max_x1_peer, %max_x1_valid = kernel.subgroup.shuffle %lane_max, %xor1, %shuffle_width : f32, i32, i32 + %max_x1 = scalar.maxnumf %lane_max, %max_x1_peer : f32 + %max_x2_peer, %max_x2_valid = kernel.subgroup.shuffle %max_x1, %xor2, %shuffle_width : f32, i32, i32 + %max_x2 = scalar.maxnumf %max_x1, %max_x2_peer : f32 + %max_x4_peer, %max_x4_valid = kernel.subgroup.shuffle %max_x2, %xor4, %shuffle_width : f32, i32, i32 + %max_x4 = scalar.maxnumf %max_x2, %max_x4_peer : f32 + %max_x8_peer, %max_x8_valid = kernel.subgroup.shuffle %max_x4, %xor8, %shuffle_width : f32, i32, i32 + %group_max = scalar.maxnumf %max_x4, %max_x8_peer : f32 + %amax = scalar.maxnumf %group_max, %amax_eps : f32 + %a_scale = scalar.mulf %amax, %inv7 : f32 + %a_rscale = scalar.divf %f7, %amax : f32 + %gid32_base = index.mul %gid, %c2 : index + %gid32_high = index.add %gid32_base, %c1 : index + %token = index.div %gid, %nchunks : index + %chunk = index.rem %gid, %nchunks : index + %plane_chunk32 = index.mul %chunk, %c2 : index + %plane_low_base = index.mul %plane_chunk32, %cols_b : index + %plane_gid32_low = index.add %plane_low_base, %token : index + %plane_chunk32_high = index.add %plane_chunk32, %c1 : index + %plane_high_base = index.mul %plane_chunk32_high, %cols_b : index + %plane_gid32_high = index.add %plane_high_base, %token : index + %scale_low_index = scf.if %plane_major -> (index) { + scf.yield %plane_gid32_low : index + } else { + scf.yield %gid32_base : index + } + %scale_high_index = scf.if %plane_major -> (index) { + scf.yield %plane_gid32_high : index + } else { + scf.yield %gid32_high : index + } + scf.if %is_scale_leader { + view.store %a_scale, %ds_view[%scale_low_index] : f32, view<[%groups]xf32> + view.store %a_scale, %ds_view[%scale_high_index] : f32, view<[%groups]xf32> + } + + %rs = vector.splat %a_rscale : vector<4xf32> + %scaled = vector.mulf %v, %rs : vector<4xf32> + %rounded = vector.roundf %scaled : vector<4xf32> + %q32 = vector.fptosi %rounded : vector<4xf32> to vector<4xi32> + %qn = vector.andi %q32, %mask15 : vector<4xi32> + %q0 = vector.extract %qn[0] : vector<4xi32> -> i32 + %q1 = vector.extract %qn[1] : vector<4xi32> -> i32 + %q2 = vector.extract %qn[2] : vector<4xi32> -> i32 + %q3 = vector.extract %qn[3] : vector<4xi32> -> i32 + %q1s = scalar.shli %q1, %shift4 : i32 + %q2s = scalar.shli %q2, %shift8 : i32 + %q3s = scalar.shli %q3, %shift12 : i32 + %q01 = scalar.ori %q0, %q1s : i32 + %q23 = scalar.ori %q2s, %q3s : i32 + %qpacked32 = scalar.ori %q01, %q23 : i32 + %qpacked = scalar.trunci %qpacked32 : i32 to i16 + %plane_chunk_base = index.mul %chunk, %cols_b : index + %plane_gid = index.add %plane_chunk_base, %token : index + %payload_gid = scf.if %plane_major -> (index) { + scf.yield %plane_gid : index + } else { + scf.yield %gid : index + } + %word_base = index.mul %payload_gid, %c16 : index + %word_index0 = index.add %word_base, %lane_in_group : index + %word_index = index.assume %word_index0 [range(%word_index0, 0, 268435455)] : index + view.store %qpacked, %qs_view[%word_index] : i16, view<[%qs_halfwords]xi16> + + %lane_sum = vector.reduce %rounded, %c0_f32 : vector<4xf32>, f32 + %sum_x1_peer, %sum_x1_valid = kernel.subgroup.shuffle %lane_sum, %xor1, %shuffle_width : f32, i32, i32 + %sum_x1 = scalar.addf %lane_sum, %sum_x1_peer : f32 + %sum_x2_peer, %sum_x2_valid = kernel.subgroup.shuffle %sum_x1, %xor2, %shuffle_width : f32, i32, i32 + %sum_x2 = scalar.addf %sum_x1, %sum_x2_peer : f32 + %sum_x4_peer, %sum_x4_valid = kernel.subgroup.shuffle %sum_x2, %xor4, %shuffle_width : f32, i32, i32 + %group_sum = scalar.addf %sum_x2, %sum_x4_peer : f32 + scf.if %is_sum_leader { + %sum_value = scalar.fptosi %group_sum : f32 to i32 + %row_sum_index = index.add %gid32_base, %group_half : index + %plane_sum_chunk = index.add %plane_chunk32, %group_half : index + %plane_sum_base = index.mul %plane_sum_chunk, %cols_b : index + %plane_sum_index = index.add %plane_sum_base, %token : index + %sum_index = scf.if %plane_major -> (index) { + scf.yield %plane_sum_index : index + } else { + scf.yield %row_sum_index : index + } + view.store %sum_value, %sum_view[%sum_index] : i32, view<[%groups]xi32> + } + } + } + template.return +} + +template.decl @ggml.mul_mat.symmetric_i4.m128n128_wg256.body(%apply_swiglu: i1, %publish_f32: i1, %publish_f16: i1, %publish_u4: i1, %gate_na: buffer, %src0_na: buffer, %dst_na: buffer, %aq_na: buffer, %as_na: buffer, %asum_na: buffer, %qout_qs_na: buffer, %qout_ds_na: buffer, %qout_sums_na: buffer, %qout_scratch: buffer, %base: offset, %qact_scale_base: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index) + +config.decl @ggml.mul_mat.symmetric_i4.input_size : %value: index where [range(%value, 32, 32768), mul(%value, 32)] + +config.decl @ggml.mul_mat.symmetric_i4.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.symmetric_i4.token_count : %value: index where [range(%value, 1, 32768)] + +// Shared Q8_0/IU8 WMMA contraction schedule. QAct storage layout is +// normalized by the two roots before this compile-time template is applied. +template.def<@ggml.mul_mat.symmetric_i4.m128n128_wg256.body> device @ggml_mul_mat_symmetric_i4_m128n128_wg256_body(%apply_swiglu: i1, %publish_f32: i1, %publish_f16: i1, %publish_u4: i1, %gate_na: buffer, %src0_na: buffer, %dst_na: buffer, %aq_na: buffer, %as_na: buffer, %asum_na: buffer, %qout_qs_na: buffer, %qout_ds_na: buffer, %qout_sums_na: buffer, %qout_scratch: buffer, %base: offset, %qact_scale_base: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index) { + %src0_i8 = buffer.view %src0_na[%base] : buffer -> view<1073741824xi8> + %src0_f16 = buffer.view %src0_na[%base] : buffer -> view<536870912xf16> + %aq_flat = buffer.view %aq_na[%base] : buffer -> view<1073741824xi8> + %as_flat = buffer.view %as_na[%qact_scale_base] : buffer -> view<33554432xf32> + %asum_flat = buffer.view %asum_na[%base] : buffer -> view<33554432xi32> + %c0_i32 = scalar.constant 0 : i32 + %pf_zero_v16i8 = vector.constant 0 : vector<16xi8> + %pf_payload_bytes = index.constant 32 : index + %q4_k_block = index.constant 256 : index + %q4_block_bytes = index.constant 144 : index + %q4_k48 = index.constant 48 : index + // dst is [rows, cols] with rows contiguous, and C is [cols, rows]. + %dst_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + %dst_f16_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf16> + %gate_view = buffer.view %gate_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + + // Packed I4 operands use 32 payload bytes per logical K64 row. + // A 40-byte LDS stride retains the conflict-avoiding padding of the IU8 body. + %scratch = buffer.assume.memory_space %qout_scratch : buffer + %i4_schema = encoding.define #encoding.operand : encoding + %u4_schema = encoding.define #encoding.operand : encoding + %al_flat = buffer.view %scratch[%base] : buffer -> view<5120xi8> + %w_off = index.constant 5120 : offset + %wl_flat = buffer.view %scratch[%w_off] : buffer -> view<5120xi8> + %as_off = index.constant 10240 : offset + %asl_view = buffer.view %scratch[%as_off] : buffer -> view<256xf32> + %ws_off = index.constant 11264 : offset + %wsl_view = buffer.view %scratch[%ws_off] : buffer -> view<256xf32> + %asum_off = index.constant 12288 : offset + %asuml_view = buffer.view %scratch[%asum_off] : buffer -> view<256xf32> + %wc_off = index.constant 13312 : offset + %wcl_view = buffer.view %scratch[%wc_off] : buffer -> view<256xf32> + + %col_base = index.mul %col_tile, %k_m : index + %row_base = index.mul %row_tile, %k_n : index + + // Staging shares. + %k_aper = index.constant 64 : index + %k_wper = index.constant 64 : index + %ashare = index.mul %tid, %k_aper : index + %ascol = index.div %ashare, %k_kc : index + %asoff = index.rem %ashare, %k_kc : index + %adbase0 = index.mul %ascol, %k_astride : index + %adbase = index.add %adbase0, %asoff : index + %ascol_g = index.add %col_base, %ascol : index + %asbase = index.mul %ascol_g, %k_b : index + // WG64 ILP8 staging maps the complete native-Q8 tile in explicit passes. + // Two scale passes cover 128 K32 blocks; four payload passes cover 256 halves. + // Scale staging: scalar legacy order or eight-scale WMMA-native vectors. + %k_asn = index.constant 256 : index + %k_aspan = index.constant 32 : index + %k_wsn = index.constant 128 : index + %k_lanes0 = index.constant 0 : index + %k_lanes1 = index.constant 128 : index + + // Four waves own four 32-token blocks and share one 64-row weight tile. + %wcol = index.rem %wave, %k_wm : index + %wrow = index.div %wave, %k_wm : index + %k_mspan = index.constant 32 : index + %k_nspan = index.constant 64 : index + %wm_off = index.mul %wcol, %k_mspan : index + %wn_off = index.mul %wrow, %k_nspan : index + %m_out = index.add %col_base, %wm_off : index + %n_out = index.add %row_base, %wn_off : index + // RDNA3 WMMA RHS payload lanes map to columns lane % 16. Point each + // lane at its own output row before constructing the four N fragments. + %lane = index.rem %tid, %k_wave : index + %lane_lo = index.rem %lane, %c16 : index + %lane_hi = index.div %lane, %c16 : index + %wn_lane = index.add %wn_off, %lane_lo : index + %lm0 = index.add %wm_off, %c0 : index + %gm0 = index.add %m_out, %c0 : index + %k_ma1 = index.constant 16 : index + %lm1 = index.add %wm_off, %k_ma1 : index + %gm1 = index.add %m_out, %k_ma1 : index + %ln0 = index.add %wn_lane, %c0 : index + %gn0 = index.add %n_out, %c0 : index + %k_nb1 = index.constant 16 : index + %ln1 = index.add %wn_lane, %k_nb1 : index + %gn1 = index.add %n_out, %k_nb1 : index + %k_nb2 = index.constant 32 : index + %ln2 = index.add %wn_lane, %k_nb2 : index + %gn2 = index.add %n_out, %k_nb2 : index + %k_nb3 = index.constant 48 : index + %ln3 = index.add %wn_lane, %k_nb3 : index + %gn3 = index.add %n_out, %k_nb3 : index + %wsb0_0 = index.add %ln0, %c0 : index + %wsb0 = index.mul %wsb0_0, %k_blocks : index + %wsb1_0 = index.add %ln1, %c0 : index + %wsb1 = index.mul %wsb1_0, %k_blocks : index + %wsb2_0 = index.add %ln2, %c0 : index + %wsb2 = index.mul %wsb2_0, %k_blocks : index + %wsb3_0 = index.add %ln3, %c0 : index + %wsb3 = index.mul %wsb3_0, %k_blocks : index + // Activation scale metadata is [tile][block][lane-half][register]. + %as_tile0 = index.div %lm0, %c16 : index + %asb0 = index.mul %as_tile0, %k_blocks : index + // Activation scale metadata is [tile][block][lane-half][register]. + %as_tile1 = index.div %lm1, %c16 : index + %asb1 = index.mul %as_tile1, %k_blocks : index + + %pf_init_row0 = index.add %col_base, %tid : index + %pf_init_row = index.assume %pf_init_row0 [range(%pf_init_row0, 0, 32767)] : index + %pf_init_plane_bytes = index.mul %cols_b, %pf_payload_bytes : index + %pf_init_chunk_base = index.mul %c0, %pf_init_plane_bytes : index + %pf_init_row_offset = index.mul %pf_init_row, %pf_payload_bytes : index + %pf_init_src0_0 = index.add %pf_init_chunk_base, %pf_init_row_offset : index + %pf_init_src0 = index.assume %pf_init_src0_0 [range(%pf_init_src0_0, 0, 536870896)] : index + %pf_init_active = index.cmp ult, %tid, %k_m : index + %pf_init_v0, %pf_init_v1 = scf.if %pf_init_active -> (vector<16xi8>, vector<16xi8>) { + %pf_init_l0 = vector.load %aq_flat[%pf_init_src0] : view<1073741824xi8> -> vector<16xi8> + %pf_init_src1_0 = index.add %pf_init_src0, %c16 : index + %pf_init_src1 = index.assume %pf_init_src1_0 [range(%pf_init_src1_0, 16, 536870912)] : index + %pf_init_l1 = vector.load %aq_flat[%pf_init_src1] : view<1073741824xi8> -> vector<16xi8> + scf.yield %pf_init_l0, %pf_init_l1 : vector<16xi8>, vector<16xi8> + } else { + scf.yield %pf_zero_v16i8, %pf_zero_v16i8 : vector<16xi8>, vector<16xi8> + } + %f0_0, %f0_1, %f0_2, %f0_3, %f1_0, %f1_1, %f1_2, %f1_3, %pf_unused_v0, %pf_unused_v1 = scf.for %chunk = [%c0 to %nchunks step %c1](%fc0_0 = %fzero : vector<8xf32>, %fc0_1 = %fzero : vector<8xf32>, %fc0_2 = %fzero : vector<8xf32>, %fc0_3 = %fzero : vector<8xf32>, %fc1_0 = %fzero : vector<8xf32>, %fc1_1 = %fzero : vector<8xf32>, %fc1_2 = %fzero : vector<8xf32>, %fc1_3 = %fzero : vector<8xf32>, %pf_carry_v0 = %pf_init_v0 : vector<16xi8>, %pf_carry_v1 = %pf_init_v1 : vector<16xi8>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<16xi8>, vector<16xi8>) { + %chunk_k = index.mul %chunk, %k_kc : index + %chunk_blk = index.mul %chunk, %k_blocks : index + + %pf_store_active = index.cmp ult, %tid, %k_m : index + scf.if %pf_store_active { + %pf_dst0_0 = index.mul %tid, %k_astride : index + %pf_dst0 = index.assume %pf_dst0_0 [range(%pf_dst0_0, 0, 5080)] : index + vector.store %pf_carry_v0, %al_flat[%pf_dst0] : vector<16xi8>, view<5120xi8> + %pf_dst1_0 = index.add %pf_dst0, %c16 : index + %pf_dst1 = index.assume %pf_dst1_0 [range(%pf_dst1_0, 16, 5096)] : index + vector.store %pf_carry_v1, %al_flat[%pf_dst1] : vector<16xi8>, view<5120xi8> + } + // The transformed 144-byte block stores eight f16 scales in field 0 + // and eight signed-I4 K32 payloads in fields 1..8. Adjacent scale pairs + // are equal because this candidate quantizes both weights and activations + // with one K64 scale. + %q4_weight_lane_active = index.cmp ult, %tid, %k_n : index + scf.if %q4_weight_lane_active { + %q4_weight_tid0 = index.rem %tid, %k_n : index + %q4_weight_tid = index.assume %q4_weight_tid0 [range(%q4_weight_tid0, 0, 127)] : index + %q4_block_count = index.div %k_b, %q4_k_block : index + %q4_block = index.div %chunk, %c4 : index + %q4_pair = index.rem %chunk, %c4 : index + %q4_row0 = index.add %row_base, %q4_weight_tid : index + %q4_row = index.assume %q4_row0 [range(%q4_row0, 0, 262143)] : index + %q4_row_group_size = index.constant 64 : index + %q4_field_count = index.constant 9 : index + %q4_field_bytes = index.constant 16 : index + %q4_row_group = index.div %q4_row, %q4_row_group_size : index + %q4_row_lane = index.rem %q4_row, %q4_row_group_size : index + %q4_group_block0 = index.mul %q4_row_group, %q4_block_count : index + %q4_group_block = index.add %q4_group_block0, %q4_block : index + %q4_group_field_base = index.mul %q4_group_block, %q4_field_count : index + %q4_header_lane_base = index.mul %q4_group_field_base, %q4_row_group_size : index + %q4_header_lane = index.add %q4_header_lane_base, %q4_row_lane : index + %q4_header_byte_index = index.mul %q4_header_lane, %q4_field_bytes : index + %q4_scale_byte_stride = index.constant 4 : index + %q4_scale_byte_offset = index.mul %q4_pair, %q4_scale_byte_stride : index + %q4_scale_byte_index = index.add %q4_header_byte_index, %q4_scale_byte_offset : index + %q4_scale_byte_base = index.cast %q4_scale_byte_index : index to offset + %q4_scale_view = buffer.view %src0_na[%q4_scale_byte_base] : buffer -> view<2xf16> + %q4_scale_pair = vector.load %q4_scale_view[%c0] : view<2xf16> -> vector<2xf16> + %q4_scale_low_f16 = vector.extract %q4_scale_pair[0] : vector<2xf16> -> f16 + %q4_scale_high_f16 = vector.extract %q4_scale_pair[1] : vector<2xf16> -> f16 + %q4_d0 = scalar.extf %q4_scale_low_f16 : f16 to f32 + %q4_d1 = scalar.extf %q4_scale_high_f16 : f16 to f32 + %q4_low_group = index.mul %q4_pair, %c2 : index + %q4_high_group = index.add %q4_low_group, %c1 : index + %q4_low_field0 = index.add %q4_low_group, %c1 : index + %q4_high_field0 = index.add %q4_high_group, %c1 : index + %q4_low_field = index.add %q4_group_field_base, %q4_low_field0 : index + %q4_high_field = index.add %q4_group_field_base, %q4_high_field0 : index + %q4_low_lane_base = index.mul %q4_low_field, %q4_row_group_size : index + %q4_high_lane_base = index.mul %q4_high_field, %q4_row_group_size : index + %q4_low_lane = index.add %q4_low_lane_base, %q4_row_lane : index + %q4_high_lane = index.add %q4_high_lane_base, %q4_row_lane : index + %q4_low_byte_index = index.mul %q4_low_lane, %q4_field_bytes : index + %q4_high_byte_index = index.mul %q4_high_lane, %q4_field_bytes : index + %q4_low_byte_base = index.cast %q4_low_byte_index : index to offset + %q4_high_byte_base = index.cast %q4_high_byte_index : index to offset + %q4_low_view = buffer.view %src0_na[%q4_low_byte_base] : buffer -> view<4xi32> + %q4_high_view = buffer.view %src0_na[%q4_high_byte_base] : buffer -> view<4xi32> + %q4_low_words = vector.load %q4_low_view[%c0] : view<4xi32> -> vector<4xi32> + %q4_high_words = vector.load %q4_high_view[%c0] : view<4xi32> -> vector<4xi32> + %q4_low_packed = vector.bitcast %q4_low_words : vector<4xi32> to vector<16xi8> + %q4_high_packed = vector.bitcast %q4_high_words : vector<4xi32> to vector<16xi8> + %q4_weight_row_base = index.mul %q4_weight_tid, %k_wstride : index + %q4_w0 = index.assume %q4_weight_row_base [range(%q4_weight_row_base, 0, 5080)] : index + %q4_w1_0 = index.add %q4_weight_row_base, %c16 : index + %q4_w1 = index.assume %q4_w1_0 [range(%q4_w1_0, 16, 5096)] : index + vector.store %q4_low_packed, %wl_flat[%q4_w0] : vector<16xi8>, view<5120xi8> + vector.store %q4_high_packed, %wl_flat[%q4_w1] : vector<16xi8>, view<5120xi8> + %q4_meta0_0 = index.mul %q4_weight_tid, %c2 : index + %q4_meta0 = index.assume %q4_meta0_0 [range(%q4_meta0_0, 0, 254)] : index + %q4_meta1_0 = index.add %q4_meta0, %c1 : index + %q4_meta1 = index.assume %q4_meta1_0 [range(%q4_meta1_0, 1, 255)] : index + view.store %q4_d0, %wsl_view[%q4_meta0] : f32, view<256xf32> + view.store %q4_d1, %wsl_view[%q4_meta1] : f32, view<256xf32> + } + // Scale planes: 256 activation and 128 weight entries. + %sca0_slot0 = index.add %tid, %k_lanes0 : index + %sca0_slot = index.assume %sca0_slot0 [range(%sca0_slot0, 0, 255)] : index + %scin_a0 = index.cmp ult, %tid, %k_lanes1 : index + scf.if %scin_a0 { + %sc0_col = index.div %sca0_slot, %k_blocks : index + %sc0_blk = index.rem %sca0_slot, %k_blocks : index + %sca0_c = index.add %col_base, %sc0_col : index + %sca0_0 = index.add %chunk_blk, %sc0_blk : index + %sca0_1 = index.mul %sca0_0, %cols_b : index + %sca0_2 = index.add %sca0_1, %sca0_c : index + %sca0 = index.assume %sca0_2 [range(%sca0_2, 0, 33554431)] : index + %scav0 = view.load %as_flat[%sca0] : view<33554432xf32> -> f32 + %sca0_tile = index.div %sc0_col, %c16 : index + %sca0_inner = index.rem %sc0_col, %c16 : index + %sca0_par = index.rem %sca0_inner, %c2 : index + %sca0_v = index.div %sca0_inner, %c2 : index + %sca0_dst0 = index.mul %sca0_tile, %k_blocks : index + %sca0_dst1 = index.add %sca0_dst0, %sc0_blk : index + %sca0_dst2 = index.mul %sca0_dst1, %c2 : index + %sca0_dst3 = index.add %sca0_dst2, %sca0_par : index + %sca0_dst4 = index.mul %sca0_dst3, %c8 : index + %sca0_dst5 = index.add %sca0_dst4, %sca0_v : index + %sca0_dst = index.assume %sca0_dst5 [range(%sca0_dst5, 0, 255)] : index + view.store %scav0, %asl_view[%sca0_dst] : f32, view<256xf32> + } + %sca1_slot0 = index.add %tid, %k_lanes1 : index + %sca1_slot = index.assume %sca1_slot0 [range(%sca1_slot0, 128, 383)] : index + %scin_a1 = index.cmp ult, %sca1_slot, %k_asn : index + scf.if %scin_a1 { + %sc1_col = index.div %sca1_slot, %k_blocks : index + %sc1_blk = index.rem %sca1_slot, %k_blocks : index + %sca1_c = index.add %col_base, %sc1_col : index + %sca1_0 = index.add %chunk_blk, %sc1_blk : index + %sca1_1 = index.mul %sca1_0, %cols_b : index + %sca1_2 = index.add %sca1_1, %sca1_c : index + %sca1 = index.assume %sca1_2 [range(%sca1_2, 0, 33554431)] : index + %scav1 = view.load %as_flat[%sca1] : view<33554432xf32> -> f32 + %sca1_tile = index.div %sc1_col, %c16 : index + %sca1_inner = index.rem %sc1_col, %c16 : index + %sca1_par = index.rem %sca1_inner, %c2 : index + %sca1_v = index.div %sca1_inner, %c2 : index + %sca1_dst0 = index.mul %sca1_tile, %k_blocks : index + %sca1_dst1 = index.add %sca1_dst0, %sc1_blk : index + %sca1_dst2 = index.mul %sca1_dst1, %c2 : index + %sca1_dst3 = index.add %sca1_dst2, %sca1_par : index + %sca1_dst4 = index.mul %sca1_dst3, %c8 : index + %sca1_dst5 = index.add %sca1_dst4, %sca1_v : index + %sca1_dst = index.assume %sca1_dst5 [range(%sca1_dst5, 0, 255)] : index + view.store %scav1, %asl_view[%sca1_dst] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %pf_next0 = index.add %chunk, %c1 : index + %pf_has_next = index.cmp ult, %pf_next0, %nchunks : index + %pf_next_chunk = scf.if %pf_has_next -> (index) { + scf.yield %pf_next0 : index + } else { + scf.yield %chunk : index + } + + %pf_next_row0 = index.add %col_base, %tid : index + %pf_next_row = index.assume %pf_next_row0 [range(%pf_next_row0, 0, 32767)] : index + %pf_next_plane_bytes = index.mul %cols_b, %pf_payload_bytes : index + %pf_next_chunk_base = index.mul %pf_next_chunk, %pf_next_plane_bytes : index + %pf_next_row_offset = index.mul %pf_next_row, %pf_payload_bytes : index + %pf_next_src0_0 = index.add %pf_next_chunk_base, %pf_next_row_offset : index + %pf_next_src0 = index.assume %pf_next_src0_0 [range(%pf_next_src0_0, 0, 536870896)] : index + %pf_next_active = index.cmp ult, %tid, %k_m : index + %pf_next_v0, %pf_next_v1 = scf.if %pf_next_active -> (vector<16xi8>, vector<16xi8>) { + %pf_next_l0 = vector.load %aq_flat[%pf_next_src0] : view<1073741824xi8> -> vector<16xi8> + %pf_next_src1_0 = index.add %pf_next_src0, %c16 : index + %pf_next_src1 = index.assume %pf_next_src1_0 [range(%pf_next_src1_0, 16, 536870912)] : index + %pf_next_l1 = vector.load %aq_flat[%pf_next_src1] : view<1073741824xi8> -> vector<16xi8> + scf.yield %pf_next_l0, %pf_next_l1 : vector<16xi8>, vector<16xi8> + } else { + scf.yield %pf_zero_v16i8, %pf_zero_v16i8 : vector<16xi8>, vector<16xi8> + } + + // Chain each pair of K32 I4 products directly in I32. This keeps + // only one matrix result live before the shared K64 scale is applied. + %blk_k_u0 = index.mul %c0, %c16 : index + %blk_k16_u0 = index.add %blk_k_u0, %c8 : index + %iz_u0 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %wsi0_2_u0 = index.add %wsb0, %c0 : index + %wsi0_u0 = index.assume %wsi0_2_u0 [range(%wsi0_2_u0, 0, 255)] : index + %wsc0_u0 = view.load %wsl_view[%wsi0_u0] : view<256xf32> -> f32 + %wsv0_u0 = vector.splat %wsc0_u0 : vector<8xf32> + %wsi1_2_u0 = index.add %wsb1, %c0 : index + %wsi1_u0 = index.assume %wsi1_2_u0 [range(%wsi1_2_u0, 0, 255)] : index + %wsc1_u0 = view.load %wsl_view[%wsi1_u0] : view<256xf32> -> f32 + %wsv1_u0 = vector.splat %wsc1_u0 : vector<8xf32> + %wsi2_2_u0 = index.add %wsb2, %c0 : index + %wsi2_u0 = index.assume %wsi2_2_u0 [range(%wsi2_2_u0, 0, 255)] : index + %wsc2_u0 = view.load %wsl_view[%wsi2_u0] : view<256xf32> -> f32 + %wsv2_u0 = vector.splat %wsc2_u0 : vector<8xf32> + %wsi3_2_u0 = index.add %wsb3, %c0 : index + %wsi3_u0 = index.assume %wsi3_2_u0 [range(%wsi3_2_u0, 0, 255)] : index + %wsc3_u0 = view.load %wsl_view[%wsi3_u0] : view<256xf32> -> f32 + %wsv3_u0 = vector.splat %wsc3_u0 : vector<8xf32> + %blk_k_u1 = index.mul %c1, %c16 : index + %blk_k16_u1 = index.add %blk_k_u1, %c8 : index + %lf0_0_lane_row_u0 = index.add %lm0, %lane_lo : index + %lf0_0_row_u0 = index.mul %lf0_0_lane_row_u0, %k_astride : index + %lf0_0_base_u0 = index.add %lf0_0_row_u0, %blk_k_u0 : index + %lf0_0_idx_u0 = index.assume %lf0_0_base_u0 [range(%lf0_0_base_u0, 0, 5112)] : index + %lf0_0_raw_u0 = vector.load %al_flat[%lf0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf0_0_words_u0 = vector.bitcast %lf0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %lf0_0_u0 = vector.fragment %lf0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf0_1_lane_row_u0 = index.add %lm0, %lane_lo : index + %lf0_1_row_u0 = index.mul %lf0_1_lane_row_u0, %k_astride : index + %lf0_1_base_u0 = index.add %lf0_1_row_u0, %blk_k16_u0 : index + %lf0_1_idx_u0 = index.assume %lf0_1_base_u0 [range(%lf0_1_base_u0, 0, 5112)] : index + %lf0_1_raw_u0 = vector.load %al_flat[%lf0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf0_1_words_u0 = vector.bitcast %lf0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %lf0_1_u0 = vector.fragment %lf0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %asi0_0_u0 = index.add %asb0, %c0 : index + %asi0_1_u0 = index.mul %asi0_0_u0, %c2 : index + %asi0_2_u0 = index.add %asi0_1_u0, %lane_hi : index + %asi0_3_u0 = index.mul %asi0_2_u0, %c8 : index + %asi0_u0 = index.assume %asi0_3_u0 [range(%asi0_3_u0, 0, 120)] : index + %asv0_u0 = vector.load %asl_view[%asi0_u0] : view<256xf32> -> vector<8xf32> + %lf0_0_lane_row_u1 = index.add %lm0, %lane_lo : index + %lf0_0_row_u1 = index.mul %lf0_0_lane_row_u1, %k_astride : index + %lf0_0_base_u1 = index.add %lf0_0_row_u1, %blk_k_u1 : index + %lf0_0_idx_u1 = index.assume %lf0_0_base_u1 [range(%lf0_0_base_u1, 0, 5112)] : index + %lf0_0_raw_u1 = vector.load %al_flat[%lf0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf0_0_words_u1 = vector.bitcast %lf0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %lf0_0_u1 = vector.fragment %lf0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf0_1_lane_row_u1 = index.add %lm0, %lane_lo : index + %lf0_1_row_u1 = index.mul %lf0_1_lane_row_u1, %k_astride : index + %lf0_1_base_u1 = index.add %lf0_1_row_u1, %blk_k16_u1 : index + %lf0_1_idx_u1 = index.assume %lf0_1_base_u1 [range(%lf0_1_base_u1, 0, 5112)] : index + %lf0_1_raw_u1 = vector.load %al_flat[%lf0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf0_1_words_u1 = vector.bitcast %lf0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %lf0_1_u1 = vector.fragment %lf0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_0_row_u0 = index.mul %ln0, %k_wstride : index + %rfs0_0_0_base_u0 = index.add %rfs0_0_0_row_u0, %blk_k_u0 : index + %rfs0_0_0_idx_u0 = index.assume %rfs0_0_0_base_u0 [range(%rfs0_0_0_base_u0, 0, 5112)] : index + %rfs0_0_0_raw_u0 = vector.load %wl_flat[%rfs0_0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_0_0_words_u0 = vector.bitcast %rfs0_0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_0_0_u0 = vector.fragment %rfs0_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_1_row_u0 = index.mul %ln0, %k_wstride : index + %rfs0_0_1_base_u0 = index.add %rfs0_0_1_row_u0, %blk_k16_u0 : index + %rfs0_0_1_idx_u0 = index.assume %rfs0_0_1_base_u0 [range(%rfs0_0_1_base_u0, 0, 5112)] : index + %rfs0_0_1_raw_u0 = vector.load %wl_flat[%rfs0_0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_0_1_words_u0 = vector.bitcast %rfs0_0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_0_1_u0 = vector.fragment %rfs0_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_0_u0 = vector.mma %lf0_0_u0, %rfs0_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_0_u0 = vector.mma %lf0_1_u0, %rfs0_0_1_u0, %i0_0_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs0_0_0_row_u1 = index.mul %ln0, %k_wstride : index + %rfs0_0_0_base_u1 = index.add %rfs0_0_0_row_u1, %blk_k_u1 : index + %rfs0_0_0_idx_u1 = index.assume %rfs0_0_0_base_u1 [range(%rfs0_0_0_base_u1, 0, 5112)] : index + %rfs0_0_0_raw_u1 = vector.load %wl_flat[%rfs0_0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_0_0_words_u1 = vector.bitcast %rfs0_0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_0_0_u1 = vector.fragment %rfs0_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_1_row_u1 = index.mul %ln0, %k_wstride : index + %rfs0_0_1_base_u1 = index.add %rfs0_0_1_row_u1, %blk_k16_u1 : index + %rfs0_0_1_idx_u1 = index.assume %rfs0_0_1_base_u1 [range(%rfs0_0_1_base_u1, 0, 5112)] : index + %rfs0_0_1_raw_u1 = vector.load %wl_flat[%rfs0_0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_0_1_words_u1 = vector.bitcast %rfs0_0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_0_1_u1 = vector.fragment %rfs0_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_0_u1 = vector.mma %lf0_0_u1, %rfs0_0_0_u1, %i1_0_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_0_u1 = vector.mma %lf0_1_u1, %rfs0_0_1_u1, %i0_0_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_0_k64 = vector.sitofp %i1_0_0_u1 : vector<8xi32> to vector<8xf32> + %sv0_0_k64 = vector.mulf %asv0_u0, %wsv0_u0 : vector<8xf32> + %fn0_0 = vector.fmaf %ff0_0_k64, %sv0_0_k64, %fc0_0 : vector<8xf32> + scf.schedule.fence + %rfs0_1_0_row_u0 = index.mul %ln1, %k_wstride : index + %rfs0_1_0_base_u0 = index.add %rfs0_1_0_row_u0, %blk_k_u0 : index + %rfs0_1_0_idx_u0 = index.assume %rfs0_1_0_base_u0 [range(%rfs0_1_0_base_u0, 0, 5112)] : index + %rfs0_1_0_raw_u0 = vector.load %wl_flat[%rfs0_1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_1_0_words_u0 = vector.bitcast %rfs0_1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_1_0_u0 = vector.fragment %rfs0_1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_1_1_row_u0 = index.mul %ln1, %k_wstride : index + %rfs0_1_1_base_u0 = index.add %rfs0_1_1_row_u0, %blk_k16_u0 : index + %rfs0_1_1_idx_u0 = index.assume %rfs0_1_1_base_u0 [range(%rfs0_1_1_base_u0, 0, 5112)] : index + %rfs0_1_1_raw_u0 = vector.load %wl_flat[%rfs0_1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_1_1_words_u0 = vector.bitcast %rfs0_1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_1_1_u0 = vector.fragment %rfs0_1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_1_u0 = vector.mma %lf0_0_u0, %rfs0_1_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_1_u0 = vector.mma %lf0_1_u0, %rfs0_1_1_u0, %i0_0_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs0_1_0_row_u1 = index.mul %ln1, %k_wstride : index + %rfs0_1_0_base_u1 = index.add %rfs0_1_0_row_u1, %blk_k_u1 : index + %rfs0_1_0_idx_u1 = index.assume %rfs0_1_0_base_u1 [range(%rfs0_1_0_base_u1, 0, 5112)] : index + %rfs0_1_0_raw_u1 = vector.load %wl_flat[%rfs0_1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_1_0_words_u1 = vector.bitcast %rfs0_1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_1_0_u1 = vector.fragment %rfs0_1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_1_1_row_u1 = index.mul %ln1, %k_wstride : index + %rfs0_1_1_base_u1 = index.add %rfs0_1_1_row_u1, %blk_k16_u1 : index + %rfs0_1_1_idx_u1 = index.assume %rfs0_1_1_base_u1 [range(%rfs0_1_1_base_u1, 0, 5112)] : index + %rfs0_1_1_raw_u1 = vector.load %wl_flat[%rfs0_1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_1_1_words_u1 = vector.bitcast %rfs0_1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_1_1_u1 = vector.fragment %rfs0_1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_1_u1 = vector.mma %lf0_0_u1, %rfs0_1_0_u1, %i1_0_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_1_u1 = vector.mma %lf0_1_u1, %rfs0_1_1_u1, %i0_0_1_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_1_k64 = vector.sitofp %i1_0_1_u1 : vector<8xi32> to vector<8xf32> + %sv0_1_k64 = vector.mulf %asv0_u0, %wsv1_u0 : vector<8xf32> + %fn0_1 = vector.fmaf %ff0_1_k64, %sv0_1_k64, %fc0_1 : vector<8xf32> + scf.schedule.fence + %rfs0_2_0_row_u0 = index.mul %ln2, %k_wstride : index + %rfs0_2_0_base_u0 = index.add %rfs0_2_0_row_u0, %blk_k_u0 : index + %rfs0_2_0_idx_u0 = index.assume %rfs0_2_0_base_u0 [range(%rfs0_2_0_base_u0, 0, 5112)] : index + %rfs0_2_0_raw_u0 = vector.load %wl_flat[%rfs0_2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_2_0_words_u0 = vector.bitcast %rfs0_2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_2_0_u0 = vector.fragment %rfs0_2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_2_1_row_u0 = index.mul %ln2, %k_wstride : index + %rfs0_2_1_base_u0 = index.add %rfs0_2_1_row_u0, %blk_k16_u0 : index + %rfs0_2_1_idx_u0 = index.assume %rfs0_2_1_base_u0 [range(%rfs0_2_1_base_u0, 0, 5112)] : index + %rfs0_2_1_raw_u0 = vector.load %wl_flat[%rfs0_2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_2_1_words_u0 = vector.bitcast %rfs0_2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_2_1_u0 = vector.fragment %rfs0_2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_2_u0 = vector.mma %lf0_0_u0, %rfs0_2_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_2_u0 = vector.mma %lf0_1_u0, %rfs0_2_1_u0, %i0_0_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs0_2_0_row_u1 = index.mul %ln2, %k_wstride : index + %rfs0_2_0_base_u1 = index.add %rfs0_2_0_row_u1, %blk_k_u1 : index + %rfs0_2_0_idx_u1 = index.assume %rfs0_2_0_base_u1 [range(%rfs0_2_0_base_u1, 0, 5112)] : index + %rfs0_2_0_raw_u1 = vector.load %wl_flat[%rfs0_2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_2_0_words_u1 = vector.bitcast %rfs0_2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_2_0_u1 = vector.fragment %rfs0_2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_2_1_row_u1 = index.mul %ln2, %k_wstride : index + %rfs0_2_1_base_u1 = index.add %rfs0_2_1_row_u1, %blk_k16_u1 : index + %rfs0_2_1_idx_u1 = index.assume %rfs0_2_1_base_u1 [range(%rfs0_2_1_base_u1, 0, 5112)] : index + %rfs0_2_1_raw_u1 = vector.load %wl_flat[%rfs0_2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_2_1_words_u1 = vector.bitcast %rfs0_2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_2_1_u1 = vector.fragment %rfs0_2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_2_u1 = vector.mma %lf0_0_u1, %rfs0_2_0_u1, %i1_0_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_2_u1 = vector.mma %lf0_1_u1, %rfs0_2_1_u1, %i0_0_2_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_2_k64 = vector.sitofp %i1_0_2_u1 : vector<8xi32> to vector<8xf32> + %sv0_2_k64 = vector.mulf %asv0_u0, %wsv2_u0 : vector<8xf32> + %fn0_2 = vector.fmaf %ff0_2_k64, %sv0_2_k64, %fc0_2 : vector<8xf32> + scf.schedule.fence + %rfs0_3_0_row_u0 = index.mul %ln3, %k_wstride : index + %rfs0_3_0_base_u0 = index.add %rfs0_3_0_row_u0, %blk_k_u0 : index + %rfs0_3_0_idx_u0 = index.assume %rfs0_3_0_base_u0 [range(%rfs0_3_0_base_u0, 0, 5112)] : index + %rfs0_3_0_raw_u0 = vector.load %wl_flat[%rfs0_3_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_3_0_words_u0 = vector.bitcast %rfs0_3_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_3_0_u0 = vector.fragment %rfs0_3_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_3_1_row_u0 = index.mul %ln3, %k_wstride : index + %rfs0_3_1_base_u0 = index.add %rfs0_3_1_row_u0, %blk_k16_u0 : index + %rfs0_3_1_idx_u0 = index.assume %rfs0_3_1_base_u0 [range(%rfs0_3_1_base_u0, 0, 5112)] : index + %rfs0_3_1_raw_u0 = vector.load %wl_flat[%rfs0_3_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs0_3_1_words_u0 = vector.bitcast %rfs0_3_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs0_3_1_u0 = vector.fragment %rfs0_3_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_3_u0 = vector.mma %lf0_0_u0, %rfs0_3_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_3_u0 = vector.mma %lf0_1_u0, %rfs0_3_1_u0, %i0_0_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs0_3_0_row_u1 = index.mul %ln3, %k_wstride : index + %rfs0_3_0_base_u1 = index.add %rfs0_3_0_row_u1, %blk_k_u1 : index + %rfs0_3_0_idx_u1 = index.assume %rfs0_3_0_base_u1 [range(%rfs0_3_0_base_u1, 0, 5112)] : index + %rfs0_3_0_raw_u1 = vector.load %wl_flat[%rfs0_3_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_3_0_words_u1 = vector.bitcast %rfs0_3_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_3_0_u1 = vector.fragment %rfs0_3_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_3_1_row_u1 = index.mul %ln3, %k_wstride : index + %rfs0_3_1_base_u1 = index.add %rfs0_3_1_row_u1, %blk_k16_u1 : index + %rfs0_3_1_idx_u1 = index.assume %rfs0_3_1_base_u1 [range(%rfs0_3_1_base_u1, 0, 5112)] : index + %rfs0_3_1_raw_u1 = vector.load %wl_flat[%rfs0_3_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs0_3_1_words_u1 = vector.bitcast %rfs0_3_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs0_3_1_u1 = vector.fragment %rfs0_3_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_3_u1 = vector.mma %lf0_0_u1, %rfs0_3_0_u1, %i1_0_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_3_u1 = vector.mma %lf0_1_u1, %rfs0_3_1_u1, %i0_0_3_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_3_k64 = vector.sitofp %i1_0_3_u1 : vector<8xi32> to vector<8xf32> + %sv0_3_k64 = vector.mulf %asv0_u0, %wsv3_u0 : vector<8xf32> + %fn0_3 = vector.fmaf %ff0_3_k64, %sv0_3_k64, %fc0_3 : vector<8xf32> + scf.schedule.fence + %lf1_0_lane_row_u0 = index.add %lm1, %lane_lo : index + %lf1_0_row_u0 = index.mul %lf1_0_lane_row_u0, %k_astride : index + %lf1_0_base_u0 = index.add %lf1_0_row_u0, %blk_k_u0 : index + %lf1_0_idx_u0 = index.assume %lf1_0_base_u0 [range(%lf1_0_base_u0, 0, 5112)] : index + %lf1_0_raw_u0 = vector.load %al_flat[%lf1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf1_0_words_u0 = vector.bitcast %lf1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %lf1_0_u0 = vector.fragment %lf1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf1_1_lane_row_u0 = index.add %lm1, %lane_lo : index + %lf1_1_row_u0 = index.mul %lf1_1_lane_row_u0, %k_astride : index + %lf1_1_base_u0 = index.add %lf1_1_row_u0, %blk_k16_u0 : index + %lf1_1_idx_u0 = index.assume %lf1_1_base_u0 [range(%lf1_1_base_u0, 0, 5112)] : index + %lf1_1_raw_u0 = vector.load %al_flat[%lf1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %lf1_1_words_u0 = vector.bitcast %lf1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %lf1_1_u0 = vector.fragment %lf1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %asi1_0_u0 = index.add %asb1, %c0 : index + %asi1_1_u0 = index.mul %asi1_0_u0, %c2 : index + %asi1_2_u0 = index.add %asi1_1_u0, %lane_hi : index + %asi1_3_u0 = index.mul %asi1_2_u0, %c8 : index + %asi1_u0 = index.assume %asi1_3_u0 [range(%asi1_3_u0, 0, 120)] : index + %asv1_u0 = vector.load %asl_view[%asi1_u0] : view<256xf32> -> vector<8xf32> + %lf1_0_lane_row_u1 = index.add %lm1, %lane_lo : index + %lf1_0_row_u1 = index.mul %lf1_0_lane_row_u1, %k_astride : index + %lf1_0_base_u1 = index.add %lf1_0_row_u1, %blk_k_u1 : index + %lf1_0_idx_u1 = index.assume %lf1_0_base_u1 [range(%lf1_0_base_u1, 0, 5112)] : index + %lf1_0_raw_u1 = vector.load %al_flat[%lf1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf1_0_words_u1 = vector.bitcast %lf1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %lf1_0_u1 = vector.fragment %lf1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf1_1_lane_row_u1 = index.add %lm1, %lane_lo : index + %lf1_1_row_u1 = index.mul %lf1_1_lane_row_u1, %k_astride : index + %lf1_1_base_u1 = index.add %lf1_1_row_u1, %blk_k16_u1 : index + %lf1_1_idx_u1 = index.assume %lf1_1_base_u1 [range(%lf1_1_base_u1, 0, 5112)] : index + %lf1_1_raw_u1 = vector.load %al_flat[%lf1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %lf1_1_words_u1 = vector.bitcast %lf1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %lf1_1_u1 = vector.fragment %lf1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_0_0_row_u0 = index.mul %ln0, %k_wstride : index + %rfs1_0_0_base_u0 = index.add %rfs1_0_0_row_u0, %blk_k_u0 : index + %rfs1_0_0_idx_u0 = index.assume %rfs1_0_0_base_u0 [range(%rfs1_0_0_base_u0, 0, 5112)] : index + %rfs1_0_0_raw_u0 = vector.load %wl_flat[%rfs1_0_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_0_0_words_u0 = vector.bitcast %rfs1_0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_0_0_u0 = vector.fragment %rfs1_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_0_1_row_u0 = index.mul %ln0, %k_wstride : index + %rfs1_0_1_base_u0 = index.add %rfs1_0_1_row_u0, %blk_k16_u0 : index + %rfs1_0_1_idx_u0 = index.assume %rfs1_0_1_base_u0 [range(%rfs1_0_1_base_u0, 0, 5112)] : index + %rfs1_0_1_raw_u0 = vector.load %wl_flat[%rfs1_0_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_0_1_words_u0 = vector.bitcast %rfs1_0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_0_1_u0 = vector.fragment %rfs1_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_0_u0 = vector.mma %lf1_0_u0, %rfs1_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_0_u0 = vector.mma %lf1_1_u0, %rfs1_0_1_u0, %i0_1_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs1_0_0_row_u1 = index.mul %ln0, %k_wstride : index + %rfs1_0_0_base_u1 = index.add %rfs1_0_0_row_u1, %blk_k_u1 : index + %rfs1_0_0_idx_u1 = index.assume %rfs1_0_0_base_u1 [range(%rfs1_0_0_base_u1, 0, 5112)] : index + %rfs1_0_0_raw_u1 = vector.load %wl_flat[%rfs1_0_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_0_0_words_u1 = vector.bitcast %rfs1_0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_0_0_u1 = vector.fragment %rfs1_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_0_1_row_u1 = index.mul %ln0, %k_wstride : index + %rfs1_0_1_base_u1 = index.add %rfs1_0_1_row_u1, %blk_k16_u1 : index + %rfs1_0_1_idx_u1 = index.assume %rfs1_0_1_base_u1 [range(%rfs1_0_1_base_u1, 0, 5112)] : index + %rfs1_0_1_raw_u1 = vector.load %wl_flat[%rfs1_0_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_0_1_words_u1 = vector.bitcast %rfs1_0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_0_1_u1 = vector.fragment %rfs1_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_0_u1 = vector.mma %lf1_0_u1, %rfs1_0_0_u1, %i1_1_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_0_u1 = vector.mma %lf1_1_u1, %rfs1_0_1_u1, %i0_1_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_0_k64 = vector.sitofp %i1_1_0_u1 : vector<8xi32> to vector<8xf32> + %sv1_0_k64 = vector.mulf %asv1_u0, %wsv0_u0 : vector<8xf32> + %fn1_0 = vector.fmaf %ff1_0_k64, %sv1_0_k64, %fc1_0 : vector<8xf32> + scf.schedule.fence + %rfs1_1_0_row_u0 = index.mul %ln1, %k_wstride : index + %rfs1_1_0_base_u0 = index.add %rfs1_1_0_row_u0, %blk_k_u0 : index + %rfs1_1_0_idx_u0 = index.assume %rfs1_1_0_base_u0 [range(%rfs1_1_0_base_u0, 0, 5112)] : index + %rfs1_1_0_raw_u0 = vector.load %wl_flat[%rfs1_1_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_1_0_words_u0 = vector.bitcast %rfs1_1_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_1_0_u0 = vector.fragment %rfs1_1_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_1_1_row_u0 = index.mul %ln1, %k_wstride : index + %rfs1_1_1_base_u0 = index.add %rfs1_1_1_row_u0, %blk_k16_u0 : index + %rfs1_1_1_idx_u0 = index.assume %rfs1_1_1_base_u0 [range(%rfs1_1_1_base_u0, 0, 5112)] : index + %rfs1_1_1_raw_u0 = vector.load %wl_flat[%rfs1_1_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_1_1_words_u0 = vector.bitcast %rfs1_1_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_1_1_u0 = vector.fragment %rfs1_1_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_1_u0 = vector.mma %lf1_0_u0, %rfs1_1_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_1_u0 = vector.mma %lf1_1_u0, %rfs1_1_1_u0, %i0_1_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs1_1_0_row_u1 = index.mul %ln1, %k_wstride : index + %rfs1_1_0_base_u1 = index.add %rfs1_1_0_row_u1, %blk_k_u1 : index + %rfs1_1_0_idx_u1 = index.assume %rfs1_1_0_base_u1 [range(%rfs1_1_0_base_u1, 0, 5112)] : index + %rfs1_1_0_raw_u1 = vector.load %wl_flat[%rfs1_1_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_1_0_words_u1 = vector.bitcast %rfs1_1_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_1_0_u1 = vector.fragment %rfs1_1_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_1_1_row_u1 = index.mul %ln1, %k_wstride : index + %rfs1_1_1_base_u1 = index.add %rfs1_1_1_row_u1, %blk_k16_u1 : index + %rfs1_1_1_idx_u1 = index.assume %rfs1_1_1_base_u1 [range(%rfs1_1_1_base_u1, 0, 5112)] : index + %rfs1_1_1_raw_u1 = vector.load %wl_flat[%rfs1_1_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_1_1_words_u1 = vector.bitcast %rfs1_1_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_1_1_u1 = vector.fragment %rfs1_1_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_1_u1 = vector.mma %lf1_0_u1, %rfs1_1_0_u1, %i1_1_1_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_1_u1 = vector.mma %lf1_1_u1, %rfs1_1_1_u1, %i0_1_1_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_1_k64 = vector.sitofp %i1_1_1_u1 : vector<8xi32> to vector<8xf32> + %sv1_1_k64 = vector.mulf %asv1_u0, %wsv1_u0 : vector<8xf32> + %fn1_1 = vector.fmaf %ff1_1_k64, %sv1_1_k64, %fc1_1 : vector<8xf32> + scf.schedule.fence + %rfs1_2_0_row_u0 = index.mul %ln2, %k_wstride : index + %rfs1_2_0_base_u0 = index.add %rfs1_2_0_row_u0, %blk_k_u0 : index + %rfs1_2_0_idx_u0 = index.assume %rfs1_2_0_base_u0 [range(%rfs1_2_0_base_u0, 0, 5112)] : index + %rfs1_2_0_raw_u0 = vector.load %wl_flat[%rfs1_2_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_2_0_words_u0 = vector.bitcast %rfs1_2_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_2_0_u0 = vector.fragment %rfs1_2_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_2_1_row_u0 = index.mul %ln2, %k_wstride : index + %rfs1_2_1_base_u0 = index.add %rfs1_2_1_row_u0, %blk_k16_u0 : index + %rfs1_2_1_idx_u0 = index.assume %rfs1_2_1_base_u0 [range(%rfs1_2_1_base_u0, 0, 5112)] : index + %rfs1_2_1_raw_u0 = vector.load %wl_flat[%rfs1_2_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_2_1_words_u0 = vector.bitcast %rfs1_2_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_2_1_u0 = vector.fragment %rfs1_2_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_2_u0 = vector.mma %lf1_0_u0, %rfs1_2_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_2_u0 = vector.mma %lf1_1_u0, %rfs1_2_1_u0, %i0_1_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs1_2_0_row_u1 = index.mul %ln2, %k_wstride : index + %rfs1_2_0_base_u1 = index.add %rfs1_2_0_row_u1, %blk_k_u1 : index + %rfs1_2_0_idx_u1 = index.assume %rfs1_2_0_base_u1 [range(%rfs1_2_0_base_u1, 0, 5112)] : index + %rfs1_2_0_raw_u1 = vector.load %wl_flat[%rfs1_2_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_2_0_words_u1 = vector.bitcast %rfs1_2_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_2_0_u1 = vector.fragment %rfs1_2_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_2_1_row_u1 = index.mul %ln2, %k_wstride : index + %rfs1_2_1_base_u1 = index.add %rfs1_2_1_row_u1, %blk_k16_u1 : index + %rfs1_2_1_idx_u1 = index.assume %rfs1_2_1_base_u1 [range(%rfs1_2_1_base_u1, 0, 5112)] : index + %rfs1_2_1_raw_u1 = vector.load %wl_flat[%rfs1_2_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_2_1_words_u1 = vector.bitcast %rfs1_2_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_2_1_u1 = vector.fragment %rfs1_2_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_2_u1 = vector.mma %lf1_0_u1, %rfs1_2_0_u1, %i1_1_2_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_2_u1 = vector.mma %lf1_1_u1, %rfs1_2_1_u1, %i0_1_2_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_2_k64 = vector.sitofp %i1_1_2_u1 : vector<8xi32> to vector<8xf32> + %sv1_2_k64 = vector.mulf %asv1_u0, %wsv2_u0 : vector<8xf32> + %fn1_2 = vector.fmaf %ff1_2_k64, %sv1_2_k64, %fc1_2 : vector<8xf32> + scf.schedule.fence + %rfs1_3_0_row_u0 = index.mul %ln3, %k_wstride : index + %rfs1_3_0_base_u0 = index.add %rfs1_3_0_row_u0, %blk_k_u0 : index + %rfs1_3_0_idx_u0 = index.assume %rfs1_3_0_base_u0 [range(%rfs1_3_0_base_u0, 0, 5112)] : index + %rfs1_3_0_raw_u0 = vector.load %wl_flat[%rfs1_3_0_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_3_0_words_u0 = vector.bitcast %rfs1_3_0_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_3_0_u0 = vector.fragment %rfs1_3_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_3_1_row_u0 = index.mul %ln3, %k_wstride : index + %rfs1_3_1_base_u0 = index.add %rfs1_3_1_row_u0, %blk_k16_u0 : index + %rfs1_3_1_idx_u0 = index.assume %rfs1_3_1_base_u0 [range(%rfs1_3_1_base_u0, 0, 5112)] : index + %rfs1_3_1_raw_u0 = vector.load %wl_flat[%rfs1_3_1_idx_u0] : view<5120xi8> -> vector<8xi8> + %rfs1_3_1_words_u0 = vector.bitcast %rfs1_3_1_raw_u0 : vector<8xi8> to vector<2xi32> + %rfs1_3_1_u0 = vector.fragment %rfs1_3_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_3_u0 = vector.mma %lf1_0_u0, %rfs1_3_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_3_u0 = vector.mma %lf1_1_u0, %rfs1_3_1_u0, %i0_1_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %rfs1_3_0_row_u1 = index.mul %ln3, %k_wstride : index + %rfs1_3_0_base_u1 = index.add %rfs1_3_0_row_u1, %blk_k_u1 : index + %rfs1_3_0_idx_u1 = index.assume %rfs1_3_0_base_u1 [range(%rfs1_3_0_base_u1, 0, 5112)] : index + %rfs1_3_0_raw_u1 = vector.load %wl_flat[%rfs1_3_0_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_3_0_words_u1 = vector.bitcast %rfs1_3_0_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_3_0_u1 = vector.fragment %rfs1_3_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs1_3_1_row_u1 = index.mul %ln3, %k_wstride : index + %rfs1_3_1_base_u1 = index.add %rfs1_3_1_row_u1, %blk_k16_u1 : index + %rfs1_3_1_idx_u1 = index.assume %rfs1_3_1_base_u1 [range(%rfs1_3_1_base_u1, 0, 5112)] : index + %rfs1_3_1_raw_u1 = vector.load %wl_flat[%rfs1_3_1_idx_u1] : view<5120xi8> -> vector<8xi8> + %rfs1_3_1_words_u1 = vector.bitcast %rfs1_3_1_raw_u1 : vector<8xi8> to vector<2xi32> + %rfs1_3_1_u1 = vector.fragment %rfs1_3_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_1_3_u1 = vector.mma %lf1_0_u1, %rfs1_3_0_u1, %i1_1_3_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_1_3_u1 = vector.mma %lf1_1_u1, %rfs1_3_1_u1, %i0_1_3_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff1_3_k64 = vector.sitofp %i1_1_3_u1 : vector<8xi32> to vector<8xf32> + %sv1_3_k64 = vector.mulf %asv1_u0, %wsv3_u0 : vector<8xf32> + %fn1_3 = vector.fmaf %ff1_3_k64, %sv1_3_k64, %fc1_3 : vector<8xf32> + scf.schedule.fence + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %fn0_0, %fn0_1, %fn0_2, %fn0_3, %fn1_0, %fn1_1, %fn1_2, %fn1_3, %pf_next_v0, %pf_next_v1 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<16xi8>, vector<16xi8> + } + + %out0_0 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_0 = vector.fragment.load %gate_view[%gm0, %gn0] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_0 = vector.siluf %gate0_0 : vector<8xf32> + %result0_0 = vector.mulf %silu0_0, %f0_0 : vector<8xf32> + scf.yield %result0_0 : vector<8xf32> + } else { + scf.yield %f0_0 : vector<8xf32> + } + %out0_1 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_1 = vector.fragment.load %gate_view[%gm0, %gn1] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_1 = vector.siluf %gate0_1 : vector<8xf32> + %result0_1 = vector.mulf %silu0_1, %f0_1 : vector<8xf32> + scf.yield %result0_1 : vector<8xf32> + } else { + scf.yield %f0_1 : vector<8xf32> + } + %out0_2 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_2 = vector.fragment.load %gate_view[%gm0, %gn2] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_2 = vector.siluf %gate0_2 : vector<8xf32> + %result0_2 = vector.mulf %silu0_2, %f0_2 : vector<8xf32> + scf.yield %result0_2 : vector<8xf32> + } else { + scf.yield %f0_2 : vector<8xf32> + } + %out0_3 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_3 = vector.fragment.load %gate_view[%gm0, %gn3] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_3 = vector.siluf %gate0_3 : vector<8xf32> + %result0_3 = vector.mulf %silu0_3, %f0_3 : vector<8xf32> + scf.yield %result0_3 : vector<8xf32> + } else { + scf.yield %f0_3 : vector<8xf32> + } + %out1_0 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_0 = vector.fragment.load %gate_view[%gm1, %gn0] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_0 = vector.siluf %gate1_0 : vector<8xf32> + %result1_0 = vector.mulf %silu1_0, %f1_0 : vector<8xf32> + scf.yield %result1_0 : vector<8xf32> + } else { + scf.yield %f1_0 : vector<8xf32> + } + %out1_1 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_1 = vector.fragment.load %gate_view[%gm1, %gn1] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_1 = vector.siluf %gate1_1 : vector<8xf32> + %result1_1 = vector.mulf %silu1_1, %f1_1 : vector<8xf32> + scf.yield %result1_1 : vector<8xf32> + } else { + scf.yield %f1_1 : vector<8xf32> + } + %out1_2 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_2 = vector.fragment.load %gate_view[%gm1, %gn2] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_2 = vector.siluf %gate1_2 : vector<8xf32> + %result1_2 = vector.mulf %silu1_2, %f1_2 : vector<8xf32> + scf.yield %result1_2 : vector<8xf32> + } else { + scf.yield %f1_2 : vector<8xf32> + } + %out1_3 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_3 = vector.fragment.load %gate_view[%gm1, %gn3] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_3 = vector.siluf %gate1_3 : vector<8xf32> + %result1_3 = vector.mulf %silu1_3, %f1_3 : vector<8xf32> + scf.yield %result1_3 : vector<8xf32> + } else { + scf.yield %f1_3 : vector<8xf32> + } + + scf.if %publish_f32 { + vector.fragment.store %out0_0, %dst_view[%gm0, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out0_1, %dst_view[%gm0, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out0_2, %dst_view[%gm0, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out0_3, %dst_view[%gm0, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_0, %dst_view[%gm1, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_1, %dst_view[%gm1, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_2, %dst_view[%gm1, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_3, %dst_view[%gm1, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + } + + scf.if %publish_f16 { + %f16_tile = buffer.view %qout_scratch[%base] : buffer -> view<128x32xf32> + %f16_c4 = index.constant 4 : index + %f16_c8 = index.constant 8 : index + %f16_c32 = index.constant 32 : index + %f16_c512 = index.constant 512 : index + scf.for %f16_phase = [%c0 to %c2 step %c1] { + %f16_first_half = index.cmp eq, %f16_phase, %c0 : index + scf.if %f16_first_half { + vector.fragment.store %out0_0, %f16_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out0_1, %f16_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_0, %f16_tile[%lm1, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_1, %f16_tile[%lm1, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + } else { + vector.fragment.store %out0_2, %f16_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out0_3, %f16_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_2, %f16_tile[%lm1, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_3, %f16_tile[%lm1, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.for %f16_batch = [%c0 to %f16_c8 step %c1] { + %f16_batch_base = index.mul %f16_batch, %f16_c512 : index + %f16_thread_base = index.mul %tid, %f16_c4 : index + %f16_linear0 = index.add %f16_batch_base, %f16_thread_base : index + %f16_linear = index.assume %f16_linear0 [range(%f16_linear0, 0, 4092), mul(%f16_linear0, 4)] : index + %f16_token_local = index.div %f16_linear, %f16_c32 : index + %f16_row_local = index.rem %f16_linear, %f16_c32 : index + %f16_wide = vector.load %f16_tile[%f16_token_local, %f16_row_local] : view<128x32xf32> -> vector<4xf32> + %f16_narrow = vector.fptrunc %f16_wide : vector<4xf32> to vector<4xf16> + %f16_token0 = index.add %col_base, %f16_token_local : index + %f16_phase_row = index.mul %f16_phase, %f16_c32 : index + %f16_row0 = index.add %row_base, %f16_phase_row : index + %f16_row1 = index.add %f16_row0, %f16_row_local : index + %f16_row_end0 = index.add %f16_row1, %f16_c4 : index + %f16_token, %f16_row, %f16_row_end, %f16_cols, %f16_rows = index.assume %f16_token0, %f16_row1, %f16_row_end0, %cols_b, %rows_b [lt(%f16_token0, %cols_b), le(%f16_row_end0, %rows_b)] : index, index, index, index, index + vector.store %f16_narrow, %dst_f16_view[%f16_token, %f16_row] : vector<4xf16>, view<[%cols_b]x[%rows_b]xf16> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + + // The following down projection consumes affine-U4 K32 blocks. Publish + // that representation directly while each 128-token x 64-row SwiGLU tile + // is live, avoiding both an F32 round trip and a standalone quantizer. + scf.if %publish_u4 { + %u4_tile = buffer.view %qout_scratch[%base] : buffer -> view<128x32xf32> + + %u4_c4 = index.constant 4 : index + %u4_c8 = index.constant 8 : index + %u4_c16 = index.constant 16 : index + %u4_c32 = index.constant 32 : index + %u4_zero_f32 = scalar.constant 0.0 : f32 + %u4_zero_i32 = scalar.constant 0 : i32 + %u4_eps = scalar.constant 1.0000000000000001e-30 : f32 + %u4_f15 = scalar.constant 15.0 : f32 + %u4_vzero = vector.constant 0.0 : vector<4xf32> + %u4_v15 = vector.constant 15.0 : vector<4xf32> + %u4_bias_neg8 = vector.constant -8 : vector<4xi32> + %u4_nibble_mask = vector.constant 15 : vector<4xi32> + %u4_c8_i32 = scalar.constant 8 : i32 + %u4_xor1 = scalar.constant 1 : i32 + %u4_xor2 = scalar.constant 2 : i32 + %u4_xor4 = scalar.constant 4 : i32 + %u4_shift4 = scalar.constant 4 : i32 + %u4_shift8 = scalar.constant 8 : i32 + %u4_shift12 = scalar.constant 12 : i32 + %u4_shuffle_width = scalar.constant 32 : i32 + + %u4_total = index.mul %rows_b, %cols_b : index + %u4_groups0 = index.div %u4_total, %u4_c32 : index + %u4_groups = index.assume %u4_groups0 [range(%u4_groups0, 1, 33554432)] : index + %u4_halfwords = index.div %u4_total, %u4_c4 : index + %u4_qs = buffer.view %qout_qs_na[%base] : buffer -> view<[%u4_halfwords]xi16> + %u4_ds = buffer.view %qout_ds_na[%base] : buffer -> view<[%u4_groups]xf32> + %u4_meta_count = index.mul %u4_groups, %c2 : index + %u4_meta = buffer.view %qout_sums_na[%base] : buffer -> view<[%u4_meta_count]xi32> + + %u4_group_in_wave = index.div %lane, %u4_c8 : index + %u4_lane_in_group = index.rem %lane, %u4_c8 : index + %u4_is_leader = index.cmp eq, %u4_lane_in_group, %c0 : index + %u4_wave_group_base = index.mul %wave, %c4 : index + %u4_local_group_base = index.add %u4_wave_group_base, %u4_group_in_wave : index + %u4_row_group_base = index.mul %row_tile, %c2 : index + %u4_row_groups = index.div %rows_b, %u4_c32 : index + + scf.for %u4_phase = [%c0 to %c2 step %c1] { + %u4_phase0 = index.cmp eq, %u4_phase, %c0 : index + scf.if %u4_phase0 { + vector.fragment.store %out0_0, %u4_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out0_1, %u4_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_0, %u4_tile[%lm1, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_1, %u4_tile[%lm1, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + } else { + vector.fragment.store %out0_2, %u4_tile[%lm0, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out0_3, %u4_tile[%lm0, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_2, %u4_tile[%lm1, %c0] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + vector.fragment.store %out1_3, %u4_tile[%lm1, %c16] shape [%c16, %c16] : vector<8xf32>, view<128x32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + scf.for %u4_batch = [%c0 to %u4_c8 step %c1] { + %u4_batch_group = index.mul %u4_batch, %u4_c16 : index + %u4_local_group0 = index.add %u4_batch_group, %u4_local_group_base : index + %u4_token_local = index.assume %u4_local_group0 [range(%u4_local_group0, 0, 127)] : index + %u4_lane_row0 = index.mul %u4_lane_in_group, %u4_c4 : index + %u4_lane_row = index.assume %u4_lane_row0 [range(%u4_lane_row0, 0, 28)] : index + %u4_values = vector.load %u4_tile[%u4_token_local, %u4_lane_row] : view<128x32xf32> -> vector<4xf32> + + %u4_lane_max = vector.reduce %u4_values, %u4_zero_f32 : vector<4xf32>, f32 + %u4_neg_values = vector.subf %u4_vzero, %u4_values : vector<4xf32> + %u4_lane_neg_min = vector.reduce %u4_neg_values, %u4_zero_f32 : vector<4xf32>, f32 + %u4_max1_peer, %u4_max1_valid = kernel.subgroup.shuffle %u4_lane_max, %u4_xor1, %u4_shuffle_width : f32, i32, i32 + %u4_max1 = scalar.maxnumf %u4_lane_max, %u4_max1_peer : f32 + %u4_max2_peer, %u4_max2_valid = kernel.subgroup.shuffle %u4_max1, %u4_xor2, %u4_shuffle_width : f32, i32, i32 + %u4_max2 = scalar.maxnumf %u4_max1, %u4_max2_peer : f32 + %u4_max4_peer, %u4_max4_valid = kernel.subgroup.shuffle %u4_max2, %u4_xor4, %u4_shuffle_width : f32, i32, i32 + %u4_group_max = scalar.maxnumf %u4_max2, %u4_max4_peer : f32 + %u4_min1_peer, %u4_min1_valid = kernel.subgroup.shuffle %u4_lane_neg_min, %u4_xor1, %u4_shuffle_width : f32, i32, i32 + %u4_min1 = scalar.maxnumf %u4_lane_neg_min, %u4_min1_peer : f32 + %u4_min2_peer, %u4_min2_valid = kernel.subgroup.shuffle %u4_min1, %u4_xor2, %u4_shuffle_width : f32, i32, i32 + %u4_min2 = scalar.maxnumf %u4_min1, %u4_min2_peer : f32 + %u4_min4_peer, %u4_min4_valid = kernel.subgroup.shuffle %u4_min2, %u4_xor4, %u4_shuffle_width : f32, i32, i32 + %u4_group_neg_min = scalar.maxnumf %u4_min2, %u4_min4_peer : f32 + %u4_range0 = scalar.addf %u4_group_max, %u4_group_neg_min : f32 + %u4_range = scalar.maxnumf %u4_range0, %u4_eps : f32 + %u4_scale = scalar.divf %u4_range, %u4_f15 : f32 + %u4_rscale = scalar.divf %u4_f15, %u4_range : f32 + %u4_zp_raw = scalar.mulf %u4_group_neg_min, %u4_rscale : f32 + %u4_zp_round = scalar.roundf %u4_zp_raw : f32 + %u4_zp_clamped = scalar.clampf %u4_zp_round, %u4_zero_f32, %u4_f15 : f32 + %u4_zp = scalar.fptosi %u4_zp_clamped : f32 to i32 + + %u4_token = index.add %col_base, %u4_token_local : index + %u4_row_group = index.add %u4_row_group_base, %u4_phase : index + %u4_gid_base = index.mul %u4_token, %u4_row_groups : index + %u4_gid0 = index.add %u4_gid_base, %u4_row_group : index + %u4_gid = index.assume %u4_gid0 [range(%u4_gid0, 0, 33554431)] : index + scf.if %u4_is_leader { + view.store %u4_scale, %u4_ds[%u4_gid] : f32, view<[%u4_groups]xf32> + } + + %u4_rs = vector.splat %u4_rscale : vector<4xf32> + %u4_zpv = vector.splat %u4_zp_clamped : vector<4xf32> + %u4_scaled = vector.mulf %u4_values, %u4_rs : vector<4xf32> + %u4_shifted = vector.addf %u4_scaled, %u4_zpv : vector<4xf32> + %u4_rounded = vector.roundf %u4_shifted : vector<4xf32> + %u4_nonnegative = vector.maxnumf %u4_rounded, %u4_vzero : vector<4xf32> + %u4_to_ceiling = vector.subf %u4_v15, %u4_nonnegative : vector<4xf32> + %u4_ceiling_nonnegative = vector.maxnumf %u4_to_ceiling, %u4_vzero : vector<4xf32> + %u4_clamped = vector.subf %u4_v15, %u4_ceiling_nonnegative : vector<4xf32> + %u4_q32_unsigned = vector.fptosi %u4_clamped : vector<4xf32> to vector<4xi32> + %u4_q32 = vector.addi %u4_q32_unsigned, %u4_bias_neg8 : vector<4xi32> + %u4_qbits = vector.andi %u4_q32, %u4_nibble_mask : vector<4xi32> + %u4_q0 = vector.extract %u4_qbits[0] : vector<4xi32> -> i32 + %u4_q1 = vector.extract %u4_qbits[1] : vector<4xi32> -> i32 + %u4_q2 = vector.extract %u4_qbits[2] : vector<4xi32> -> i32 + %u4_q3 = vector.extract %u4_qbits[3] : vector<4xi32> -> i32 + %u4_q1s = scalar.shli %u4_q1, %u4_shift4 : i32 + %u4_q2s = scalar.shli %u4_q2, %u4_shift8 : i32 + %u4_q3s = scalar.shli %u4_q3, %u4_shift12 : i32 + %u4_q01 = scalar.ori %u4_q0, %u4_q1s : i32 + %u4_q23 = scalar.ori %u4_q2s, %u4_q3s : i32 + %u4_packed32 = scalar.ori %u4_q01, %u4_q23 : i32 + %u4_packed = scalar.trunci %u4_packed32 : i32 to i16 + %u4_word_base = index.mul %u4_gid, %u4_c8 : index + %u4_word_index0 = index.add %u4_word_base, %u4_lane_in_group : index + %u4_word_index = index.assume %u4_word_index0 [range(%u4_word_index0, 0, 268435455)] : index + view.store %u4_packed, %u4_qs[%u4_word_index] : i16, view<[%u4_halfwords]xi16> + + %u4_lane_sum = vector.reduce %u4_q32, %u4_zero_i32 : vector<4xi32>, i32 + %u4_sum1_peer, %u4_sum1_valid = kernel.subgroup.shuffle %u4_lane_sum, %u4_xor1, %u4_shuffle_width : i32, i32, i32 + %u4_sum1 = scalar.addi %u4_lane_sum, %u4_sum1_peer : i32 + %u4_sum2_peer, %u4_sum2_valid = kernel.subgroup.shuffle %u4_sum1, %u4_xor2, %u4_shuffle_width : i32, i32, i32 + %u4_sum2 = scalar.addi %u4_sum1, %u4_sum2_peer : i32 + %u4_sum4_peer, %u4_sum4_valid = kernel.subgroup.shuffle %u4_sum2, %u4_xor4, %u4_shuffle_width : i32, i32, i32 + %u4_group_sum = scalar.addi %u4_sum2, %u4_sum4_peer : i32 + scf.if %u4_is_leader { + %u4_zp_signed = scalar.subi %u4_zp, %u4_c8_i32 : i32 + %u4_meta0 = index.mul %u4_gid, %c2 : index + %u4_meta1 = index.add %u4_meta0, %c1 : index + view.store %u4_group_sum, %u4_meta[%u4_meta0] : i32, view<[%u4_meta_count]xi32> + view.store %u4_zp_signed, %u4_meta[%u4_meta1] : i32, view<[%u4_meta_count]xi32> + } + } + kernel.barrier scope(workgroup) ordering(acq_rel) + } + } + template.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_wmma") @ggml_mul_mat_symmetric_i4_wmma() { + %unit = index.constant 1 : index + %k_n = index.constant 128 : index + %k_m = index.constant 128 : index + %n_m1 = index.constant 127 : index + %m_m1 = index.constant 127 : index + %wg = index.constant 256 : index + %rows = config.get @ggml.mul_mat.symmetric_i4.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.token_count : index + %rows_up = index.add %rows, %n_m1 : index + %row_tiles = index.div %rows_up, %k_n : index + %cols_up = index.add %cols, %m_m1 : index + %col_tiles = index.div %cols_up, %k_m : index + %workgroups = index.mul %col_tiles, %row_tiles : index + kernel.launch.config workgroups(%workgroups, %unit, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%src0: buffer, %src1: buffer, %dst: buffer, %aq: buffer, %as: buffer, %asum: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 40 : index + %k_wstride = index.constant 40 : index + %k_m = index.constant 128 : index + %k_n = index.constant 128 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 4 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %apply_swiglu = scalar.constant false : i1 + %publish_f32 = scalar.constant true : i1 + %publish_f16 = scalar.constant false : i1 + %publish_u4 = scalar.constant false : i1 + %work_scratch_bytes = index.constant 14336 : offset + %work_scratch = buffer.alloca align(16) %work_scratch_bytes : buffer + + %k = config.get @ggml.mul_mat.symmetric_i4.input_size : index + %rows = config.get @ggml.mul_mat.symmetric_i4.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.token_count : index + %k_b = index.assume %k [range(%k, 64, 32768)] : index + %rows_b = index.assume %rows [range(%rows, 1, 262144)] : index + %cols_b = index.assume %cols [range(%cols, 1, 32768)] : index + %nchunks = index.div %k_b, %k_kc : index + %kblocks = index.div %k_b, %c32 : index + + // Keep each 16-row-tile weight footprint inside the 32 MiB Infinity Cache, + // while preserving immediate reuse across all token tiles in the block. + %grid_c127 = index.constant 127 : index + %grid_block_rows = index.constant 16 : index + %grid_row_tiles_up = index.add %rows_b, %grid_c127 : index + %grid_row_tiles = index.div %grid_row_tiles_up, %k_n : index + %grid_col_tiles_up = index.add %cols_b, %grid_c127 : index + %grid_col_tiles = index.div %grid_col_tiles_up, %k_m : index + %grid_full_blocks = index.div %grid_row_tiles, %grid_block_rows : index + %grid_full_rows = index.mul %grid_full_blocks, %grid_block_rows : index + %grid_full_groups = index.mul %grid_full_rows, %grid_col_tiles : index + %grid_tail_rows = index.sub %grid_row_tiles, %grid_full_rows : index + %grid_block_span = index.mul %grid_block_rows, %grid_col_tiles : index + %grid_wgid0 = kernel.workgroup.id : index + %grid_is_full = index.cmp ult, %grid_wgid0, %grid_full_groups : index + %col_tile0, %row_tile0 = scf.if %grid_is_full -> (index, index) { + %grid_block = index.div %grid_wgid0, %grid_block_span : index + %grid_within = index.rem %grid_wgid0, %grid_block_span : index + %grid_col = index.div %grid_within, %grid_block_rows : index + %grid_row_in_block = index.rem %grid_within, %grid_block_rows : index + %grid_row_base = index.mul %grid_block, %grid_block_rows : index + %grid_row = index.add %grid_row_base, %grid_row_in_block : index + scf.yield %grid_col, %grid_row : index, index + } else { + %grid_tail_wgid = index.sub %grid_wgid0, %grid_full_groups : index + %grid_col = index.div %grid_tail_wgid, %grid_tail_rows : index + %grid_row_in_tail = index.rem %grid_tail_wgid, %grid_tail_rows : index + %grid_row = index.add %grid_full_rows, %grid_row_in_tail : index + scf.yield %grid_col, %grid_row : index, index + } + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 511)] : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 4095)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %wave = index.div %tid, %k_wave : index + + %src0_g = buffer.assume.memory_space %src0 : buffer + %dst_g = buffer.assume.memory_space %dst : buffer + %aq_g = buffer.assume.memory_space %aq : buffer + %as_g = buffer.assume.memory_space %as : buffer + %asum_g = buffer.assume.memory_space %asum : buffer + %src0_na, %dst_na, %aq_na, %as_na, %asum_na = buffer.assume.noalias %src0_g, %dst_g, %aq_g, %as_g, %asum_g : buffer, buffer, buffer, buffer, buffer + + template.apply<@ggml.mul_mat.symmetric_i4.m128n128_wg256.body>(%apply_swiglu, %publish_f32, %publish_f16, %publish_u4, %dst, %src0_na, %dst_na, %aq_na, %as_na, %asum_na, %dst_na, %dst_na, %dst_na, %work_scratch, %base, %base, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %k_b, %rows_b, %cols_b, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave) : (i1, i1, i1, i1, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index) + kernel.return +} + +check.case public @ggml_mul_mat_symmetric_i4_wmma_k5120_n6144_m512_case { + %weight = check.generate.fill value(17) : tensor<17694720xi8> + %input = check.generate.fill value(0.125) : tensor<2621440xf32> + %output = check.generate.fill value(0.0) : tensor<3145728xf32> + %qact = check.generate.fill value(17) : tensor<1310720xi8> + %scale = check.generate.fill value(0.03125) : tensor<81920xf32> + %sum = check.generate.fill value(32) : tensor<81920xi32> + kernel.launch @ggml_mul_mat_symmetric_i4_wmma(%weight, %input, %output, %qact, %scale, %sum) : (tensor<17694720xi8>, tensor<2621440xf32>, tensor<3145728xf32>, tensor<1310720xi8>, tensor<81920xf32>, tensor<81920xi32>) + check.return +} + +check.benchmark<@ggml_mul_mat_symmetric_i4_wmma_k5120_n6144_m512_case> @ggml_mul_mat_symmetric_i4_wmma_k5120_n6144_m512 + +check.case public @ggml_mul_mat_symmetric_i4_wmma_k5120_n10240_m512_case { + %weight = check.generate.fill value(17) : tensor<29491200xi8> + %input = check.generate.fill value(0.125) : tensor<2621440xf32> + %output = check.generate.fill value(0.0) : tensor<5242880xf32> + %qact = check.generate.fill value(17) : tensor<1310720xi8> + %scale = check.generate.fill value(0.03125) : tensor<81920xf32> + %sum = check.generate.fill value(32) : tensor<81920xi32> + kernel.launch @ggml_mul_mat_symmetric_i4_wmma(%weight, %input, %output, %qact, %scale, %sum) : (tensor<29491200xi8>, tensor<2621440xf32>, tensor<5242880xf32>, tensor<1310720xi8>, tensor<81920xf32>, tensor<81920xi32>) + check.return +} + +check.benchmark<@ggml_mul_mat_symmetric_i4_wmma_k5120_n10240_m512_case> @ggml_mul_mat_symmetric_i4_wmma_k5120_n10240_m512 + +template.decl @ggml.mul_mat.symmetric_i4.lowrow.m16n16.body(%scratch: buffer, %src0_na: buffer, %dst_na: buffer, %aq_na: buffer, %as_na: buffer, %asum_na: buffer, %compact_q2: i1, %split_k: i1, %partial_na: buffer, %completion_counters_na: buffer, %base: offset, %qact_scale_base: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index, %partition: index, %chunk_begin: index, %chunk_end: index) + +// Capacity-neutral symmetric-I4 weights with shared K64 scales. +// The dual-projection roots share one 128-token activation tile. +config.decl @ggml.mul_mat.symmetric_i4.lowrow.input_size : %value: index where [range(%value, 64, 32768), mul(%value, 64)] + +config.decl @ggml.mul_mat.symmetric_i4.lowrow.output_size : %value: index where [range(%value, 1, 262144)] + +config.decl @ggml.mul_mat.symmetric_i4.lowrow.token_count : %value: index where [range(%value, 1, 16)] + +config.decl @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : %value: index where [range(%value, 32, 262144), mul(%value, 32)] + +template.def<@ggml.mul_mat.symmetric_i4.lowrow.m16n16.body> device @ggml_mul_mat_symmetric_i4_lowrow_m16n16_body(%scratch: buffer, %src0_na: buffer, %dst_na: buffer, %aq_na: buffer, %as_na: buffer, %asum_na: buffer, %compact_q2: i1, %split_k: i1, %partial_na: buffer, %completion_counters_na: buffer, %base: offset, %qact_scale_base: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index, %partition: index, %chunk_begin: index, %chunk_end: index) { + %src0_i8 = buffer.view %src0_na[%base] : buffer -> view<1073741824xi8> + %src0_f16 = buffer.view %src0_na[%base] : buffer -> view<536870912xf16> + %aq_flat = buffer.view %aq_na[%base] : buffer -> view<1073741824xi8> + %as_flat = buffer.view %as_na[%qact_scale_base] : buffer -> view<33554432xf32> + %asum_flat = buffer.view %asum_na[%base] : buffer -> view<67108864xi32> + %c0_i32 = scalar.constant 0 : i32 + %q4_k_block = index.constant 256 : index + %q4_block_bytes = index.constant 144 : index + %q4_k48 = index.constant 48 : index + %q4_shift4v = vector.constant 4 : vector<4xi32> + %q4_shift8v = vector.constant 8 : vector<4xi32> + %q4_shift16v = vector.constant 16 : vector<4xi32> + %q4_nibble_mask = vector.constant 252645135 : vector<4xi32> + %q4_pair_mask = vector.constant 16711935 : vector<4xi32> + %q4_half_mask = vector.constant 65535 : vector<4xi32> + %q4_f32_32 = scalar.constant 32.0 : f32 + %c15 = index.constant 15 : index + // dst is [rows, cols] with rows contiguous, and C is [cols, rows]. + %dst_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + + // Packed I4 operands use 32 payload bytes per logical K64 row. + // A 40-byte LDS stride retains the conflict-avoiding padding of the IU8 body. + %i4_schema = encoding.define #encoding.operand : encoding + %al_flat = buffer.view %scratch[%base] : buffer -> view<640xi8> + %as_off = index.constant 640 : offset + %asl_view = buffer.view %scratch[%as_off] : buffer -> view<32xf32> + %azp_off = index.constant 768 : offset + %azpl_view = buffer.view %scratch[%azp_off] : buffer -> view<32xf32> + %asum_off = index.constant 896 : offset + %asuml_view = buffer.view %scratch[%asum_off] : buffer -> view<32xf32> + %result_off = index.constant 1024 : offset + %result_wave_bytes = index.constant 1024 : offset + %result_wave_offset = index.scale %wave, %result_wave_bytes : index, offset -> offset + %result_offset = index.add %result_off, %result_wave_offset : offset + %result_fragment_view = buffer.view %scratch[%result_offset] : buffer -> view<16x16xf32> + %result_physical_view = buffer.view %scratch[%result_offset] : buffer -> view<16x16xf32> + %arrival_scratch_view = buffer.view %scratch[%result_off] : buffer -> view<1xi32> + + %col_base = index.mul %col_tile, %k_m : index + %row_base = index.mul %row_tile, %k_n : index + + %k_aper = index.constant 64 : index + %k_wper = index.constant 64 : index + %ashare = index.mul %tid, %k_aper : index + %ascol = index.div %ashare, %k_kc : index + %asoff = index.rem %ashare, %k_kc : index + %adbase0 = index.mul %ascol, %k_astride : index + %adbase = index.add %adbase0, %asoff : index + %ascol_g = index.add %col_base, %ascol : index + %asbase = index.mul %ascol_g, %k_b : index + // WG64 ILP8 stages 128 scales in two passes and 256 payload halves in four passes. + %k_asn = index.mul %cols_b, %k_blocks : index + %k_aspan = index.constant 32 : index + %k_wsn = index.constant 128 : index + %k_lanes0 = index.constant 0 : index + %k_lanes1 = index.constant 128 : index + + // All waves share the padded M16 token tile; each owns one disjoint N16 weight slice. + %wm_off = index.constant 0 : index + %k_nspan = index.constant 16 : index + %wn_off = index.mul %wave, %k_nspan : index + %m_out = index.add %col_base, %wm_off : index + %n_out = index.add %row_base, %wn_off : index + // RDNA3 WMMA RHS lanes map to column lane % 16, so each lane starts at its output row. + %lane = index.rem %tid, %k_wave : index + %lane_lo = index.rem %lane, %c16 : index + %lane_hi = index.div %lane, %c16 : index + %wn_lane = index.add %wn_off, %lane_lo : index + %lm0 = index.add %wm_off, %c0 : index + %gm0 = index.add %m_out, %c0 : index + %ln0 = index.add %wn_lane, %c0 : index + %gn0 = index.add %n_out, %c0 : index + %wsb0_0 = index.add %ln0, %c0 : index + %wsb0 = index.mul %wsb0_0, %k_blocks : index + // Activation scale metadata is [tile][block][lane-half][register]. + %as_tile0 = index.div %lm0, %c16 : index + %asb0 = index.mul %as_tile0, %k_blocks : index + + // Carry one prefetched K64 activation payload between loop iterations. + %q4a_lane_active = index.cmp ult, %tid, %cols_b : index + %q4a_row0 = index.add %col_base, %tid : index + %q4a_row = index.assume %q4a_row0 [range(%q4a_row0, 0, 32767)] : index + %q4a_row_bytes = index.div %k_b, %c2 : index + %q4a_row_base = index.mul %q4a_row, %q4a_row_bytes : index + %q4a_c16 = index.constant 16 : index + %q4a_zero = vector.constant 0 : vector<16xi8> + %q4a_initial_v0, %q4a_initial_v1 = scf.if %q4a_lane_active -> (vector<16xi8>, vector<16xi8>) { + %q4a_initial_chunk_k = index.mul %chunk_begin, %k_kc : index + %q4a_initial_chunk_bytes = index.div %q4a_initial_chunk_k, %c2 : index + %q4a_initial_src0_0 = index.add %q4a_row_base, %q4a_initial_chunk_bytes : index + %q4a_initial_src0 = index.assume %q4a_initial_src0_0 [range(%q4a_initial_src0_0, 0, 536870896)] : index + %q4a_loaded_initial_v0 = vector.load %aq_flat[%q4a_initial_src0] : view<1073741824xi8> -> vector<16xi8> + %q4a_initial_src1_0 = index.add %q4a_initial_src0, %q4a_c16 : index + %q4a_initial_src1 = index.assume %q4a_initial_src1_0 [range(%q4a_initial_src1_0, 16, 536870912)] : index + %q4a_loaded_initial_v1 = vector.load %aq_flat[%q4a_initial_src1] : view<1073741824xi8> -> vector<16xi8> + scf.yield %q4a_loaded_initial_v0, %q4a_loaded_initial_v1 : vector<16xi8>, vector<16xi8> + } else { + scf.yield %q4a_zero, %q4a_zero : vector<16xi8>, vector<16xi8> + } + + %f0_0, %q4a_unused_v0, %q4a_unused_v1 = scf.for %chunk = [%chunk_begin to %chunk_end step %c1](%fc0_0 = %fzero : vector<8xf32>, %q4a_current_v0 = %q4a_initial_v0 : vector<16xi8>, %q4a_current_v1 = %q4a_initial_v1 : vector<16xi8>) -> (vector<8xf32>, vector<16xi8>, vector<16xi8>) { + %chunk_k = index.mul %chunk, %k_kc : index + %chunk_blk = index.mul %chunk, %k_blocks : index + + // The first 128 lanes stage the unchanged M128 activation tile. + scf.if %q4a_lane_active { + %q4a_dst0_0 = index.mul %tid, %k_astride : index + %q4a_dst0 = index.assume %q4a_dst0_0 [range(%q4a_dst0_0, 0, 600)] : index + vector.store %q4a_current_v0, %al_flat[%q4a_dst0] : vector<16xi8>, view<640xi8> + %q4a_dst1_0 = index.add %q4a_dst0, %q4a_c16 : index + %q4a_dst1 = index.assume %q4a_dst1_0 [range(%q4a_dst1_0, 16, 616)] : index + vector.store %q4a_current_v1, %al_flat[%q4a_dst1] : vector<16xi8>, view<640xi8> + } + + // Issue the following K64 payload before weight unpacking and contraction. + // The uniform last-iteration guard prevents a speculative next-row access. + %q4a_next_chunk = index.add %chunk, %c1 : index + %q4a_has_next = index.cmp ult, %q4a_next_chunk, %chunk_end : index + %q4a_load_next = scalar.andi %q4a_has_next, %q4a_lane_active : i1 + %q4a_next_v0, %q4a_next_v1 = scf.if %q4a_load_next -> (vector<16xi8>, vector<16xi8>) { + %q4a_next_chunk_k = index.mul %q4a_next_chunk, %k_kc : index + %q4a_next_chunk_bytes = index.div %q4a_next_chunk_k, %c2 : index + %q4a_next_src0_0 = index.add %q4a_row_base, %q4a_next_chunk_bytes : index + %q4a_next_src0 = index.assume %q4a_next_src0_0 [range(%q4a_next_src0_0, 0, 536870896)] : index + %q4a_loaded_v0 = vector.load %aq_flat[%q4a_next_src0] : view<1073741824xi8> -> vector<16xi8> + %q4a_next_src1_0 = index.add %q4a_next_src0, %q4a_c16 : index + %q4a_next_src1 = index.assume %q4a_next_src1_0 [range(%q4a_next_src1_0, 16, 536870912)] : index + %q4a_loaded_v1 = vector.load %aq_flat[%q4a_next_src1] : view<1073741824xi8> -> vector<16xi8> + scf.yield %q4a_loaded_v0, %q4a_loaded_v1 : vector<16xi8>, vector<16xi8> + } else { + scf.yield %q4a_zero, %q4a_zero : vector<16xi8>, vector<16xi8> + } + scf.schedule.fence + // One half-wave owns each signed-I4 K32 group for the same + // 16 output rows. Four adjacent rows share each K32 scale. + %q4_weight_tid0 = index.rem %tid, %c16 : index + %q4_weight_tid = index.assume %q4_weight_tid0 [range(%q4_weight_tid0, 0, 15)] : index + %q4_group_half = index.div %lane, %c16 : index + %q4_block_count = index.div %k_b, %q4_k_block : index + %q4_block = index.div %chunk, %c4 : index + %q4_pair = index.rem %chunk, %c4 : index + %q4_row0 = index.add %n_out, %q4_weight_tid : index + %q4_row = index.assume %q4_row0 [range(%q4_row0, 0, 262143)] : index + %q4_row_group_size0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %q4_row_group_size = index.assume %q4_row_group_size0 [range(%q4_row_group_size0, 32, 262144), mul(%q4_row_group_size0, 32)] : index + %q4_scale_cohort_size = index.constant 4 : index + %q4_field_bytes = index.constant 16 : index + %q2_field_bytes = index.constant 8 : index + %weight_field_bytes = scf.select %compact_q2, %q2_field_bytes, %q4_field_bytes : index + %q4_scale_plane_bytes = index.mul %q4_row_group_size, %c4 : index + %q4_payload_plane_bytes = index.mul %q4_row_group_size, %weight_field_bytes : index + %q4_full_bytes_per_row = index.constant 132 : index + %q2_compact_bytes_per_row = index.constant 68 : index + %q4_bytes_per_row = scf.select %compact_q2, %q2_compact_bytes_per_row, %q4_full_bytes_per_row : index + %q4_full_payload_bytes_per_row = index.constant 128 : index + %q2_compact_payload_bytes_per_row = index.constant 64 : index + %q4_payload_bytes_per_row = scf.select %compact_q2, %q2_compact_payload_bytes_per_row, %q4_full_payload_bytes_per_row : index + %q4_group_block_bytes = index.mul %q4_row_group_size, %q4_bytes_per_row : index + %q4_payload_block_bytes = index.mul %q4_row_group_size, %q4_payload_bytes_per_row : index + %q4_row_group = index.div %q4_row, %q4_row_group_size : index + %q4_row_lane = index.rem %q4_row, %q4_row_group_size : index + %q4_row_group_span0 = index.mul %q4_block_count, %q4_group_block_bytes : index + %q4_row_group_byte_base = index.mul %q4_row_group, %q4_row_group_span0 : index + %q4_payload_block_offset = index.mul %q4_block, %q4_payload_block_bytes : index + %q4_payload_block_base = index.add %q4_row_group_byte_base, %q4_payload_block_offset : index + %q4_all_payload_bytes = index.mul %q4_block_count, %q4_payload_block_bytes : index + %q4_scale_region = index.add %q4_row_group_byte_base, %q4_all_payload_bytes : index + %q4_scale_block_offset = index.mul %q4_block, %q4_scale_plane_bytes : index + %q4_scale_block_base = index.add %q4_scale_region, %q4_scale_block_offset : index + %q4_scale_cohort = index.div %q4_row_lane, %q4_scale_cohort_size : index + %q4_header_cohort_byte = index.mul %q4_scale_cohort, %q4_field_bytes : index + %q4_header_byte_index0 = index.add %q4_scale_block_base, %q4_header_cohort_byte : index + %q4_scale_byte_stride = index.constant 4 : index + %q4_scale_byte_offset = index.mul %q4_pair, %q4_scale_byte_stride : index + %q4_scale_byte_index = index.add %q4_header_byte_index0, %q4_scale_byte_offset : index + %q4_scale_byte_base = index.cast %q4_scale_byte_index : index to offset + %q4_scale_view = buffer.view %src0_na[%q4_scale_byte_base] : buffer -> view<2xf16> + %q4_scale_pair = vector.load %q4_scale_view[%c0] : view<2xf16> -> vector<2xf16> + %q4_scale_low_f16 = vector.extract %q4_scale_pair[0] : vector<2xf16> -> f16 + %q4_scale_high_f16 = vector.extract %q4_scale_pair[1] : vector<2xf16> -> f16 + %q4_scale_low = scalar.extf %q4_scale_low_f16 : f16 to f32 + %q4_scale_high = scalar.extf %q4_scale_high_f16 : f16 to f32 + %q4_scale_is_high_half = index.cmp eq, %q4_group_half, %c1 : index + %q4_d0 = scf.select %q4_scale_is_high_half, %q4_scale_high, %q4_scale_low : f32 + %q4_group_base = index.mul %q4_pair, %c2 : index + %q4_group = index.add %q4_group_base, %q4_group_half : index + %q4_payload_group_byte = index.mul %q4_group, %q4_payload_plane_bytes : index + %q4_payload_row_byte = index.mul %q4_row_lane, %weight_field_bytes : index + %q4_payload_base1 = index.add %q4_payload_block_base, %q4_payload_group_byte : index + %q4_byte_index = index.add %q4_payload_base1, %q4_payload_row_byte : index + %q4_byte_base = index.cast %q4_byte_index : index to offset + %q4_words = scf.if %compact_q2 -> (vector<4xi32>) { + %q2_view = buffer.view %src0_na[%q4_byte_base] : buffer -> view<2xi32> + %q2_packed = vector.load %q2_view[%c0] : view<2xi32> -> vector<2xi32> + %q2_mask16 = vector.constant 65535 : vector<2xi32> + %q2_mask8 = vector.constant 16711935 : vector<2xi32> + %q2_mask4 = vector.constant 252645135 : vector<2xi32> + %q2_mask2 = vector.constant 858993459 : vector<2xi32> + %q2_sign = vector.constant 572662306 : vector<2xi32> + %q2_shift16 = vector.constant 16 : vector<2xi32> + %q2_shift8 = vector.constant 8 : vector<2xi32> + %q2_shift4 = vector.constant 4 : vector<2xi32> + %q2_shift2 = vector.constant 2 : vector<2xi32> + %q2_shift1 = vector.constant 1 : vector<2xi32> + %q2_lo0 = vector.andi %q2_packed, %q2_mask16 : vector<2xi32> + %q2_hi0 = vector.shrui %q2_packed, %q2_shift16 : vector<2xi32> + %q2_lo1s = vector.shli %q2_lo0, %q2_shift8 : vector<2xi32> + %q2_lo1o = vector.ori %q2_lo0, %q2_lo1s : vector<2xi32> + %q2_lo1 = vector.andi %q2_lo1o, %q2_mask8 : vector<2xi32> + %q2_lo2s = vector.shli %q2_lo1, %q2_shift4 : vector<2xi32> + %q2_lo2o = vector.ori %q2_lo1, %q2_lo2s : vector<2xi32> + %q2_lo2 = vector.andi %q2_lo2o, %q2_mask4 : vector<2xi32> + %q2_lo3s = vector.shli %q2_lo2, %q2_shift2 : vector<2xi32> + %q2_lo3o = vector.ori %q2_lo2, %q2_lo3s : vector<2xi32> + %q2_lo3 = vector.andi %q2_lo3o, %q2_mask2 : vector<2xi32> + %q2_losign = vector.andi %q2_lo3, %q2_sign : vector<2xi32> + %q2_losign2 = vector.shli %q2_losign, %q2_shift1 : vector<2xi32> + %q2_losign3 = vector.shli %q2_losign, %q2_shift2 : vector<2xi32> + %q2_lo4 = vector.ori %q2_lo3, %q2_losign2 : vector<2xi32> + %q2_lo = vector.ori %q2_lo4, %q2_losign3 : vector<2xi32> + %q2_hi1s = vector.shli %q2_hi0, %q2_shift8 : vector<2xi32> + %q2_hi1o = vector.ori %q2_hi0, %q2_hi1s : vector<2xi32> + %q2_hi1 = vector.andi %q2_hi1o, %q2_mask8 : vector<2xi32> + %q2_hi2s = vector.shli %q2_hi1, %q2_shift4 : vector<2xi32> + %q2_hi2o = vector.ori %q2_hi1, %q2_hi2s : vector<2xi32> + %q2_hi2 = vector.andi %q2_hi2o, %q2_mask4 : vector<2xi32> + %q2_hi3s = vector.shli %q2_hi2, %q2_shift2 : vector<2xi32> + %q2_hi3o = vector.ori %q2_hi2, %q2_hi3s : vector<2xi32> + %q2_hi3 = vector.andi %q2_hi3o, %q2_mask2 : vector<2xi32> + %q2_hisign = vector.andi %q2_hi3, %q2_sign : vector<2xi32> + %q2_hisign2 = vector.shli %q2_hisign, %q2_shift1 : vector<2xi32> + %q2_hisign3 = vector.shli %q2_hisign, %q2_shift2 : vector<2xi32> + %q2_hi4 = vector.ori %q2_hi3, %q2_hisign2 : vector<2xi32> + %q2_hi = vector.ori %q2_hi4, %q2_hisign3 : vector<2xi32> + %q2_word0 = vector.extract %q2_lo[0] : vector<2xi32> -> i32 + %q2_word1 = vector.extract %q2_hi[0] : vector<2xi32> -> i32 + %q2_word2 = vector.extract %q2_lo[1] : vector<2xi32> -> i32 + %q2_word3 = vector.extract %q2_hi[1] : vector<2xi32> -> i32 + %q2_words = vector.from_elements %q2_word0, %q2_word1, %q2_word2, %q2_word3 : vector<4xi32> + scf.yield %q2_words : vector<4xi32> + } else { + %q4_view = buffer.view %src0_na[%q4_byte_base] : buffer -> view<4xi32> + %q4_loaded = vector.load %q4_view[%c0] : view<4xi32> -> vector<4xi32> + scf.yield %q4_loaded : vector<4xi32> + } + // Scale planes: 256 activation and 128 weight entries. + %sca0_slot0 = index.add %tid, %k_lanes0 : index + %sca0_slot = index.assume %sca0_slot0 [range(%sca0_slot0, 0, 63)] : index + %scin_a0 = index.cmp ult, %sca0_slot, %k_asn : index + scf.if %scin_a0 { + %sc0_col = index.div %sca0_slot, %k_blocks : index + %sc0_blk = index.rem %sca0_slot, %k_blocks : index + %sca0_c = index.add %col_base, %sc0_col : index + %sca0_0 = index.mul %sca0_c, %kblocks : index + %sca0_1 = index.add %sca0_0, %chunk_blk : index + %sca0_2 = index.add %sca0_1, %sc0_blk : index + %sca0 = index.assume %sca0_2 [range(%sca0_2, 0, 33554431)] : index + %scav0 = view.load %as_flat[%sca0] : view<33554432xf32> -> f32 + %sca0_tile = index.div %sc0_col, %c16 : index + %sca0_inner = index.rem %sc0_col, %c16 : index + %sca0_par = index.rem %sca0_inner, %c2 : index + %sca0_v = index.div %sca0_inner, %c2 : index + %sca0_dst0 = index.mul %sca0_tile, %k_blocks : index + %sca0_dst1 = index.add %sca0_dst0, %sc0_blk : index + %sca0_dst2 = index.mul %sca0_dst1, %c2 : index + %sca0_dst3 = index.add %sca0_dst2, %sca0_par : index + %sca0_dst4 = index.mul %sca0_dst3, %c8 : index + %sca0_dst5 = index.add %sca0_dst4, %sca0_v : index + %sca0_dst = index.assume %sca0_dst5 [range(%sca0_dst5, 0, 31)] : index + view.store %scav0, %asl_view[%sca0_dst] : f32, view<32xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + + %q4_xor16 = scalar.constant 16 : i32 + %q4_shuffle_width = scalar.constant 32 : i32 + %q4_word0_local = vector.extract %q4_words[0] : vector<4xi32> -> i32 + %q4_word1_local = vector.extract %q4_words[1] : vector<4xi32> -> i32 + %q4_word2_local = vector.extract %q4_words[2] : vector<4xi32> -> i32 + %q4_word3_local = vector.extract %q4_words[3] : vector<4xi32> -> i32 + %q4_word0_peer, %q4_word0_valid = kernel.subgroup.shuffle %q4_word0_local, %q4_xor16, %q4_shuffle_width : i32, i32, i32 + %q4_word1_peer, %q4_word1_valid = kernel.subgroup.shuffle %q4_word1_local, %q4_xor16, %q4_shuffle_width : i32, i32, i32 + %q4_word2_peer, %q4_word2_valid = kernel.subgroup.shuffle %q4_word2_local, %q4_xor16, %q4_shuffle_width : i32, i32, i32 + %q4_word3_peer, %q4_word3_valid = kernel.subgroup.shuffle %q4_word3_local, %q4_xor16, %q4_shuffle_width : i32, i32, i32 + %q4_d0_peer, %q4_d0_valid = kernel.subgroup.shuffle %q4_d0, %q4_xor16, %q4_shuffle_width : f32, i32, i32 + %q4_is_high_half = index.cmp eq, %q4_group_half, %c1 : index + %q4_word0_u0 = scf.select %q4_is_high_half, %q4_word0_peer, %q4_word0_local : i32 + %q4_word1_u0 = scf.select %q4_is_high_half, %q4_word1_peer, %q4_word1_local : i32 + %q4_word2_u0 = scf.select %q4_is_high_half, %q4_word2_peer, %q4_word2_local : i32 + %q4_word3_u0 = scf.select %q4_is_high_half, %q4_word3_peer, %q4_word3_local : i32 + %q4_d_u0 = scf.select %q4_is_high_half, %q4_d0_peer, %q4_d0 : f32 + %q4_d_u1 = scf.select %q4_is_high_half, %q4_d0, %q4_d0_peer : f32 + %q4_word0_u1 = scf.select %q4_is_high_half, %q4_word0_local, %q4_word0_peer : i32 + %q4_word1_u1 = scf.select %q4_is_high_half, %q4_word1_local, %q4_word1_peer : i32 + %q4_word2_u1 = scf.select %q4_is_high_half, %q4_word2_local, %q4_word2_peer : i32 + %q4_word3_u1 = scf.select %q4_is_high_half, %q4_word3_local, %q4_word3_peer : i32 + + // Apply each signed-I4 K32 dot product with its own weight scale. + %blk_k_u0 = index.mul %c0, %c16 : index + %blk_k16_u0 = index.add %blk_k_u0, %c8 : index + %iz_u0 = vector.fragment %izero shape [%c16, %c16] : vector<8xi32> + %wsv0_u0 = vector.splat %q4_d_u0 : vector<8xf32> + %lf0_0_lane_row_u0 = index.add %lm0, %lane_lo : index + %lf0_0_row_u0 = index.mul %lf0_0_lane_row_u0, %k_astride : index + %lf0_0_base_u0 = index.add %lf0_0_row_u0, %blk_k_u0 : index + %lf0_0_idx_u0 = index.assume %lf0_0_base_u0 [range(%lf0_0_base_u0, 0, 632)] : index + %lf0_0_raw_u0 = vector.load %al_flat[%lf0_0_idx_u0] : view<640xi8> -> vector<8xi8> + %lf0_0_words_u0 = vector.bitcast %lf0_0_raw_u0 : vector<8xi8> to vector<2xi32> + %lf0_0_u0 = vector.fragment %lf0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf0_1_lane_row_u0 = index.add %lm0, %lane_lo : index + %lf0_1_row_u0 = index.mul %lf0_1_lane_row_u0, %k_astride : index + %lf0_1_base_u0 = index.add %lf0_1_row_u0, %blk_k16_u0 : index + %lf0_1_idx_u0 = index.assume %lf0_1_base_u0 [range(%lf0_1_base_u0, 0, 632)] : index + %lf0_1_raw_u0 = vector.load %al_flat[%lf0_1_idx_u0] : view<640xi8> -> vector<8xi8> + %lf0_1_words_u0 = vector.bitcast %lf0_1_raw_u0 : vector<8xi8> to vector<2xi32> + %lf0_1_u0 = vector.fragment %lf0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %asi0_0_u0 = index.add %asb0, %c0 : index + %asi0_1_u0 = index.mul %asi0_0_u0, %c2 : index + %asi0_2_u0 = index.add %asi0_1_u0, %lane_hi : index + %asi0_3_u0 = index.mul %asi0_2_u0, %c8 : index + %asi0_u0 = index.assume %asi0_3_u0 [range(%asi0_3_u0, 0, 24)] : index + %asv0_u0 = vector.load %asl_view[%asi0_u0] : view<32xf32> -> vector<8xf32> + %rfs0_0_0_words_u0 = vector.from_elements %q4_word0_u0, %q4_word1_u0 : vector<2xi32> + %rfs0_0_0_u0 = vector.fragment %rfs0_0_0_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_1_words_u0 = vector.from_elements %q4_word2_u0, %q4_word3_u0 : vector<2xi32> + %rfs0_0_1_u0 = vector.fragment %rfs0_0_1_words_u0 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_0_u0 = vector.mma %lf0_0_u0, %rfs0_0_0_u0, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_0_u0 = vector.mma %lf0_1_u0, %rfs0_0_1_u0, %i0_0_0_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_0_u0 = vector.sitofp %i1_0_0_u0 : vector<8xi32> to vector<8xf32> + %sv0_0_u0 = vector.mulf %asv0_u0, %wsv0_u0 : vector<8xf32> + %fn0_0_u0 = vector.fmaf %ff0_0_u0, %sv0_0_u0, %fc0_0 : vector<8xf32> + scf.schedule.fence + + %blk_k_u1 = index.mul %c1, %c16 : index + %blk_k16_u1 = index.add %blk_k_u1, %c8 : index + %lf0_0_lane_row_u1 = index.add %lm0, %lane_lo : index + %lf0_0_row_u1 = index.mul %lf0_0_lane_row_u1, %k_astride : index + %lf0_0_base_u1 = index.add %lf0_0_row_u1, %blk_k_u1 : index + %lf0_0_idx_u1 = index.assume %lf0_0_base_u1 [range(%lf0_0_base_u1, 0, 632)] : index + %lf0_0_raw_u1 = vector.load %al_flat[%lf0_0_idx_u1] : view<640xi8> -> vector<8xi8> + %lf0_0_words_u1 = vector.bitcast %lf0_0_raw_u1 : vector<8xi8> to vector<2xi32> + %lf0_0_u1 = vector.fragment %lf0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %lf0_1_lane_row_u1 = index.add %lm0, %lane_lo : index + %lf0_1_row_u1 = index.mul %lf0_1_lane_row_u1, %k_astride : index + %lf0_1_base_u1 = index.add %lf0_1_row_u1, %blk_k16_u1 : index + %lf0_1_idx_u1 = index.assume %lf0_1_base_u1 [range(%lf0_1_base_u1, 0, 632)] : index + %lf0_1_raw_u1 = vector.load %al_flat[%lf0_1_idx_u1] : view<640xi8> -> vector<8xi8> + %lf0_1_words_u1 = vector.bitcast %lf0_1_raw_u1 : vector<8xi8> to vector<2xi32> + %lf0_1_u1 = vector.fragment %lf0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_0_words_u1 = vector.from_elements %q4_word0_u1, %q4_word1_u1 : vector<2xi32> + %rfs0_0_0_u1 = vector.fragment %rfs0_0_0_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %rfs0_0_1_words_u1 = vector.from_elements %q4_word2_u1, %q4_word3_u1 : vector<2xi32> + %rfs0_0_1_u1 = vector.fragment %rfs0_0_1_words_u1 shape [%c16, %c16] using {schema = %i4_schema : encoding} : vector<2xi32> + %i0_0_0_u1 = vector.mma %lf0_0_u1, %rfs0_0_0_u1, %iz_u0 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %i1_0_0_u1 = vector.mma %lf0_1_u1, %rfs0_0_1_u1, %i0_0_0_u1 : vector<2xi32>, vector<2xi32>, vector<8xi32> + %ff0_0_u1 = vector.sitofp %i1_0_0_u1 : vector<8xi32> to vector<8xf32> + %asi0_0_u1 = index.add %asb0, %c1 : index + %asi0_1_u1 = index.mul %asi0_0_u1, %c2 : index + %asi0_2_u1 = index.add %asi0_1_u1, %lane_hi : index + %asi0_3_u1 = index.mul %asi0_2_u1, %c8 : index + %asi0_u1 = index.assume %asi0_3_u1 [range(%asi0_3_u1, 0, 24)] : index + %asv0_u1 = vector.load %asl_view[%asi0_u1] : view<32xf32> -> vector<8xf32> + %wsv0_u1 = vector.splat %q4_d_u1 : vector<8xf32> + %sv0_0_u1 = vector.mulf %asv0_u1, %wsv0_u1 : vector<8xf32> + %fn0_0 = vector.fmaf %ff0_0_u1, %sv0_0_u1, %fn0_0_u0 : vector<8xf32> + scf.schedule.fence + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %fn0_0, %q4a_next_v0, %q4a_next_v1 : vector<8xf32>, vector<16xi8>, vector<16xi8> + } + + vector.fragment.store %f0_0, %result_fragment_view[%c0, %c0] shape [%c16, %c16] : vector<8xf32>, view<16x16xf32> + kernel.barrier scope(subgroup) ordering(acq_rel) + %publish_two = index.constant 2 : index + %publish_eight = index.constant 8 : index + %publish_token = index.div %lane, %publish_two : index + %publish_packet = index.rem %lane, %publish_two : index + %publish_channel_add = index.mul %publish_packet, %publish_eight : index + %token = index.add %gm0, %publish_token : index + %channel = index.add %n_out, %publish_channel_add : index + %valid_token = index.cmp ult, %token, %cols_b : index + %valid_channel = index.cmp ult, %channel, %rows_b : index + %writes = scalar.andi %valid_token, %valid_channel : i1 + scf.if %split_k { + %row_tiles_up = index.add %rows_b, %c15 : index + %row_tile_count = index.div %row_tiles_up, %c16 : index + %col_tiles_up = index.add %cols_b, %c15 : index + %col_tile_count = index.div %col_tiles_up, %c16 : index + %completion_counter_count = index.mul %row_tile_count, %col_tile_count : index + %completion_counters_aligned = buffer.assume.alignment %completion_counters_na {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counters_aligned[%base] : buffer -> view<[%completion_counter_count]xi32> + %is_partition_zero = index.cmp eq, %partition, %c0 : index + scf.if %writes { + %bounded_token, %output_token_count = index.assume %token, %cols_b [lt(%token, %cols_b)] : index, index + %bounded_output_view = buffer.view %dst_na[%base] : buffer -> view<[%output_token_count]x[%rows_b]xf32> + %bounded_partial_view = buffer.view %partial_na[%base] : buffer -> view<[%output_token_count]x[%rows_b]xf32> + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<8xf32> + %mask = vector.mask.range [%channel to %rows_b step %c1] : index -> vector<8xi1> + scf.if %is_partition_zero { + vector.store.mask %values, %bounded_output_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%output_token_count]x[%rows_b]xf32>, vector<8xi1> + } else { + vector.store.mask %values, %bounded_partial_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%output_token_count]x[%rows_b]xf32>, vector<8xi1> + } + } + + // Both K partitions publish before the second arrival combines P0 then P1 and resets the tile counter. + kernel.barrier scope(workgroup) ordering(release) + %is_arrival_workitem = index.cmp eq, %tid, %c0 : index + %counter_group_base = index.mul %col_tile, %row_tile_count : index + %counter_index = index.add %counter_group_base, %row_tile : index + %c1_i32 = scalar.constant 1 : i32 + %cneg2_i32 = scalar.constant -2 : i32 + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%counter_index] {ordering = acq_rel, scope = device} : i32, view<[%completion_counter_count]xi32> -> i32 + view.store %old_counter, %arrival_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %arrival_scratch_view[%c0] : view<1xi32> -> i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %c1_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + scf.if %writes { + %bounded_token, %output_token_count = index.assume %token, %cols_b [lt(%token, %cols_b)] : index, index + %bounded_output_view = buffer.view %dst_na[%base] : buffer -> view<[%output_token_count]x[%rows_b]xf32> + %bounded_partial_view = buffer.view %partial_na[%base] : buffer -> view<[%output_token_count]x[%rows_b]xf32> + %mask = vector.mask.range [%channel to %rows_b step %c1] : index -> vector<8xi1> + %partition_zero_values = vector.load.mask %bounded_output_view[%bounded_token, %channel], %mask, %fzero : view<[%output_token_count]x[%rows_b]xf32>, vector<8xi1>, vector<8xf32> + %partition_one_values = vector.load.mask %bounded_partial_view[%bounded_token, %channel], %mask, %fzero : view<[%output_token_count]x[%rows_b]xf32>, vector<8xi1>, vector<8xf32> + %combined_values = vector.addf %partition_zero_values, %partition_one_values : vector<8xf32> + vector.store.mask %combined_values, %bounded_output_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%output_token_count]x[%rows_b]xf32>, vector<8xi1> + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %cneg2_i32, %completion_counter_view[%counter_index] {ordering = release, scope = device} : i32, view<[%completion_counter_count]xi32> + } + } + } else { + scf.if %writes { + %bounded_token, %output_token_count = index.assume %token, %cols_b [lt(%token, %cols_b)] : index, index + %bounded_output_view = buffer.view %dst_na[%base] : buffer -> view<[%output_token_count]x[%rows_b]xf32> + %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf32> -> vector<8xf32> + %mask = vector.mask.range [%channel to %rows_b step %c1] : index -> vector<8xi1> + vector.store.mask %values, %bounded_output_view[%bounded_token, %channel], %mask : vector<8xf32>, view<[%output_token_count]x[%rows_b]xf32>, vector<8xi1> + } + } + template.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_wmma") @ggml_mul_mat_symmetric_i4_lowrow_wmma() { + %unit = index.constant 1 : index + %c2 = index.constant 2 : index + %k_n = index.constant 16 : index + %k_m = index.constant 16 : index + %n_m1 = index.constant 15 : index + %m_m1 = index.constant 15 : index + %wg = index.constant 64 : index + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %rows_up = index.add %rows, %n_m1 : index + %row_tiles = index.div %rows_up, %k_n : index + %row_tiles_paired_up = index.add %row_tiles, %unit : index + %row_groups = index.div %row_tiles_paired_up, %c2 : index + %cols_up = index.add %cols, %m_m1 : index + %col_tiles = index.div %cols_up, %k_m : index + kernel.launch.config workgroups(%col_tiles, %row_groups, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %aq: buffer, %as: buffer, %asum: buffer) { + %base = index.constant 0 : offset + %lds_bytes = index.constant 3072 : offset + %scratch = buffer.alloca align(16) %lds_bytes : buffer + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 40 : index + %k_wstride = index.constant 40 : index + %k_m = index.constant 16 : index + %k_n = index.constant 16 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 1 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %no_split_k = scalar.constant false : i1 + %compact_q2 = scalar.constant false : i1 + + %k = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k_b = index.assume %k [range(%k, 64, 32768)] : index + %rows_b = index.assume %rows [range(%rows, 1, 262144)] : index + %cols_b = index.assume %cols [range(%cols, 1, 32768)] : index + %nchunks = index.div %k_b, %k_kc : index + %kblocks = index.div %k_b, %c32 : index + + %col_tile0 = kernel.workgroup.id : index + %row_group0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 511)] : index + %row_group = index.assume %row_group0 [range(%row_group0, 0, 8191)] : index + %row_tile0 = index.mul %row_group, %c2 : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 16383)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 63)] : index + %wave = index.div %tid, %k_wave : index + + %weight_g = buffer.assume.memory_space %weight : buffer + %output_g = buffer.assume.memory_space %output : buffer + %aq_g = buffer.assume.memory_space %aq : buffer + %as_g = buffer.assume.memory_space %as : buffer + %asum_g = buffer.assume.memory_space %asum : buffer + %weight_na, %output_na, %aq_na, %as_na, %asum_na = buffer.assume.noalias %weight_g, %output_g, %aq_g, %as_g, %asum_g : buffer, buffer, buffer, buffer, buffer + template.apply<@ggml.mul_mat.symmetric_i4.lowrow.m16n16.body>(%scratch, %weight_na, %output_na, %aq_na, %as_na, %asum_na, %compact_q2, %no_split_k, %output_na, %output_na, %base, %base, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %k_b, %rows_b, %cols_b, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave, %c0, %c0, %nchunks) : (buffer, buffer, buffer, buffer, buffer, buffer, i1, i1, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index, index, index, index) + kernel.return +} + +// Two adjacent projections share one launch and one quantized activation. +// Each projection keeps the same low-row contraction body and owns disjoint +// output storage; the graph matcher supplies any semantically related pair. +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma") @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma() { + %unit = index.constant 1 : index + %c2 = index.constant 2 : index + %k_n = index.constant 16 : index + %k_m = index.constant 16 : index + %n_m1 = index.constant 15 : index + %m_m1 = index.constant 15 : index + %wg = index.constant 64 : index + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %rows_up = index.add %rows, %n_m1 : index + %row_tiles = index.div %rows_up, %k_n : index + %row_tiles_paired_up = index.add %row_tiles, %unit : index + %row_groups = index.div %row_tiles_paired_up, %c2 : index + %projection_row_groups = index.mul %row_groups, %c2 : index + %cols_up = index.add %cols, %m_m1 : index + %col_tiles = index.div %cols_up, %k_m : index + kernel.launch.config workgroups(%col_tiles, %projection_row_groups, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer, %aq: buffer, %as: buffer, %asum: buffer) { + %base = index.constant 0 : offset + %lds_bytes = index.constant 3072 : offset + %scratch = buffer.alloca align(16) %lds_bytes : buffer + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 40 : index + %k_wstride = index.constant 40 : index + %k_m = index.constant 16 : index + %k_n = index.constant 16 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 1 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %no_split_k = scalar.constant false : i1 + %compact_q2 = scalar.constant false : i1 + + %k = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k_b = index.assume %k [range(%k, 64, 32768)] : index + %rows_b = index.assume %rows [range(%rows, 1, 262144)] : index + %cols_b = index.assume %cols [range(%cols, 1, 16)] : index + %nchunks = index.div %k_b, %k_kc : index + %kblocks = index.div %k_b, %c32 : index + + %col_tile0 = kernel.workgroup.id : index + %projection_row_group0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 511)] : index + %projection_row_group = index.assume %projection_row_group0 [range(%projection_row_group0, 0, 16383)] : index + %projection = index.rem %projection_row_group, %c2 : index + %row_group = index.div %projection_row_group, %c2 : index + %row_tile0 = index.mul %row_group, %c2 : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 16383)] : index + %use_second = index.cmp eq, %projection, %c1 : index + %tid = index.assume %tid0 [range(%tid0, 0, 63)] : index + %wave = index.div %tid, %k_wave : index + + %first_weight_g = buffer.assume.memory_space %first_weight : buffer + %second_weight_g = buffer.assume.memory_space %second_weight : buffer + %first_output_g = buffer.assume.memory_space %first_output : buffer + %second_output_g = buffer.assume.memory_space %second_output : buffer + %aq_g = buffer.assume.memory_space %aq : buffer + %as_g = buffer.assume.memory_space %as : buffer + %asum_g = buffer.assume.memory_space %asum : buffer + scf.if %use_second { + %weight_na, %output_na, %aq_na, %as_na, %asum_na = buffer.assume.noalias %second_weight_g, %second_output_g, %aq_g, %as_g, %asum_g : buffer, buffer, buffer, buffer, buffer + template.apply<@ggml.mul_mat.symmetric_i4.lowrow.m16n16.body>(%scratch, %weight_na, %output_na, %aq_na, %as_na, %asum_na, %compact_q2, %no_split_k, %output_na, %output_na, %base, %base, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %k_b, %rows_b, %cols_b, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave, %c0, %c0, %nchunks) : (buffer, buffer, buffer, buffer, buffer, buffer, i1, i1, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index, index, index, index) + } else { + %weight_na, %output_na, %aq_na, %as_na, %asum_na = buffer.assume.noalias %first_weight_g, %first_output_g, %aq_g, %as_g, %asum_g : buffer, buffer, buffer, buffer, buffer + template.apply<@ggml.mul_mat.symmetric_i4.lowrow.m16n16.body>(%scratch, %weight_na, %output_na, %aq_na, %as_na, %asum_na, %compact_q2, %no_split_k, %output_na, %output_na, %base, %base, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %k_b, %rows_b, %cols_b, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave, %c0, %c0, %nchunks) : (buffer, buffer, buffer, buffer, buffer, buffer, i1, i1, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index, index, index, index) + } + kernel.return +} + +// Exact low-column direct-dot specializations for expanding split-K projections. +func.def inline @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%acc: f32, %weight_words: vector<4xi32>, %weight_scale: f32, %group: index, %token: index, %k: index, %group_count: index, %qact: buffer, %scales: buffer) -> (f32) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %zero_i32 = scalar.constant 0 : i32 + %zero_i32x4 = vector.constant 0 : vector<4xi32> + %qact_flat = buffer.view %qact[%base] : buffer -> view<1073741824xi8> + %scale_flat = buffer.view %scales[%base] : buffer -> view<33554432xf32> + + %token_bytes = index.div %k, %c2 : index + %token_payload_base = index.mul %token, %token_bytes : index + %group_payload_add = index.mul %group, %c16 : index + %payload_index0 = index.add %token_payload_base, %group_payload_add : index + %payload_index = index.assume %payload_index0 [range(%payload_index0, 0, 1073741808)] : index + %activation_bytes = vector.load %qact_flat[%payload_index] : view<1073741824xi8> -> vector<16xi8> + %activation_words = vector.bitcast %activation_bytes : vector<16xi8> to vector<4xi32> + %dot_parts = vector.dot8i4 %weight_words, %activation_words, %zero_i32x4 : vector<4xi32> + %dot_i32 = vector.reduce %dot_parts, %zero_i32 : vector<4xi32>, i32 + %dot = scalar.sitofp %dot_i32 : i32 to f32 + + %token_scale_base = index.mul %token, %group_count : index + %scale_index0 = index.add %token_scale_base, %group : index + %scale_index = index.assume %scale_index0 [range(%scale_index0, 0, 33554431)] : index + %activation_scale = view.load %scale_flat[%scale_index] : view<33554432xf32> -> f32 + %combined_scale = scalar.mulf %weight_scale, %activation_scale : f32 + %next = scalar.fmaf %dot, %combined_scale, %acc : f32 + func.return %next : f32 +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c1") @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c1() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %rows_up = index.add %rows, %c127 : index + %row_workgroups = index.div %rows_up, %c128 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %qact: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 1, 1)] : index + %base = index.constant 0 : offset + %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 + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + + %split_workgroup = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %partition = index.rem %split_workgroup, %c2 : index + %workgroup = index.div %split_workgroup, %c2 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %row_base = index.mul %workgroup, %c128 : index + %row = index.add %row_base, %workitem : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %groups_per_partition = index.div %group_count, %c2 : index + %group_begin = index.mul %partition, %groups_per_partition : index + %group_end = index.add %group_begin, %groups_per_partition : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %weight_global = buffer.assume.memory_space %weight : buffer + %output_global = buffer.assume.memory_space %output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %partial_global = buffer.assume.memory_space %partial : buffer + %counters_global = buffer.assume.memory_space %completion_counters : buffer + %weight_na, %output_na, %qact_na, %scales_na, %partial_na, %counters_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global, %partial_global, %counters_global : buffer, buffer, buffer, buffer, buffer, buffer + + %result0 = scf.if %valid_row -> (f32) { + %acc0 = scf.for %group = [%group_begin to %group_end step %c1](%iter0 = %zero_f32 : f32) -> (f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index = index.add %payload_index1, %row_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<4xi32> + %weight_words = vector.load %weight_view[%c0] : view<4xi32> -> vector<4xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %next0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter0, %weight_words, %weight_scale, %group, %c0, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + scf.yield %next0 : f32 + } + scf.yield %acc0 : f32 + } else { + scf.yield %zero_f32 : f32 + } + + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %partial_view = buffer.view %partial_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %is_partition_zero = index.cmp eq, %partition, %c0 : index + scf.if %valid_row { + scf.if %is_partition_zero { + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + } else { + view.store %result0, %partial_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + } + } + + kernel.barrier scope(workgroup) ordering(release) + %counter_count = index.div %rows_plus, %c128 : index + %counter_view = buffer.view %counters_na[%base] : buffer -> view<[%counter_count]xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + %one_i32 = scalar.constant 1 : i32 + %zero_i32 = scalar.constant 0 : i32 + %negative_two_i32 = scalar.constant -2 : i32 + %local_old_counter = scf.if %is_arrival_lane -> (i32) { + %old_counter = view.atomic.rmw %one_i32, %counter_view[%workgroup] {ordering = acq_rel, scope = device} : i32, view<[%counter_count]xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %zero_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %one_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + scf.if %valid_row { + %partition_zero0 = view.load %output_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one0 = view.load %partial_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %combined0 = scalar.addf %partition_zero0, %partition_one0 : f32 + view.store %combined0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_two_i32, %counter_view[%workgroup] {ordering = release, scope = device} : i32, view<[%counter_count]xi32> + } + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c2") @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c2() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %rows_up = index.add %rows, %c127 : index + %row_workgroups = index.div %rows_up, %c128 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %qact: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 2, 2)] : index + %base = index.constant 0 : offset + %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 + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + + %split_workgroup = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %partition = index.rem %split_workgroup, %c2 : index + %workgroup = index.div %split_workgroup, %c2 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %row_base = index.mul %workgroup, %c128 : index + %row = index.add %row_base, %workitem : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %groups_per_partition = index.div %group_count, %c2 : index + %group_begin = index.mul %partition, %groups_per_partition : index + %group_end = index.add %group_begin, %groups_per_partition : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %weight_global = buffer.assume.memory_space %weight : buffer + %output_global = buffer.assume.memory_space %output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %partial_global = buffer.assume.memory_space %partial : buffer + %counters_global = buffer.assume.memory_space %completion_counters : buffer + %weight_na, %output_na, %qact_na, %scales_na, %partial_na, %counters_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global, %partial_global, %counters_global : buffer, buffer, buffer, buffer, buffer, buffer + + %result0, %result1 = scf.if %valid_row -> (f32, f32) { + %acc0, %acc1 = scf.for %group = [%group_begin to %group_end step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32) -> (f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index = index.add %payload_index1, %row_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<4xi32> + %weight_words = vector.load %weight_view[%c0] : view<4xi32> -> vector<4xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %next0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter0, %weight_words, %weight_scale, %group, %c0, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter1, %weight_words, %weight_scale, %group, %c1, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + scf.yield %next0, %next1 : f32, f32 + } + scf.yield %acc0, %acc1 : f32, f32 + } else { + scf.yield %zero_f32, %zero_f32 : f32, f32 + } + + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %partial_view = buffer.view %partial_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %is_partition_zero = index.cmp eq, %partition, %c0 : index + scf.if %valid_row { + scf.if %is_partition_zero { + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + } else { + view.store %result0, %partial_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %partial_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + } + } + + kernel.barrier scope(workgroup) ordering(release) + %counter_count = index.div %rows_plus, %c128 : index + %counter_view = buffer.view %counters_na[%base] : buffer -> view<[%counter_count]xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + %one_i32 = scalar.constant 1 : i32 + %zero_i32 = scalar.constant 0 : i32 + %negative_two_i32 = scalar.constant -2 : i32 + %local_old_counter = scf.if %is_arrival_lane -> (i32) { + %old_counter = view.atomic.rmw %one_i32, %counter_view[%workgroup] {ordering = acq_rel, scope = device} : i32, view<[%counter_count]xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %zero_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %one_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + scf.if %valid_row { + %partition_zero0 = view.load %output_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero1 = view.load %output_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one0 = view.load %partial_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one1 = view.load %partial_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %combined0 = scalar.addf %partition_zero0, %partition_one0 : f32 + %combined1 = scalar.addf %partition_zero1, %partition_one1 : f32 + view.store %combined0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_two_i32, %counter_view[%workgroup] {ordering = release, scope = device} : i32, view<[%counter_count]xi32> + } + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c3") @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c3() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %rows_up = index.add %rows, %c127 : index + %row_workgroups = index.div %rows_up, %c128 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %qact: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 3, 3)] : index + %base = index.constant 0 : offset + %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 + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + + %split_workgroup = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %partition = index.rem %split_workgroup, %c2 : index + %workgroup = index.div %split_workgroup, %c2 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %row_base = index.mul %workgroup, %c128 : index + %row = index.add %row_base, %workitem : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %groups_per_partition = index.div %group_count, %c2 : index + %group_begin = index.mul %partition, %groups_per_partition : index + %group_end = index.add %group_begin, %groups_per_partition : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %weight_global = buffer.assume.memory_space %weight : buffer + %output_global = buffer.assume.memory_space %output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %partial_global = buffer.assume.memory_space %partial : buffer + %counters_global = buffer.assume.memory_space %completion_counters : buffer + %weight_na, %output_na, %qact_na, %scales_na, %partial_na, %counters_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global, %partial_global, %counters_global : buffer, buffer, buffer, buffer, buffer, buffer + + %result0, %result1, %result2 = scf.if %valid_row -> (f32, f32, f32) { + %acc0, %acc1, %acc2 = scf.for %group = [%group_begin to %group_end step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32) -> (f32, f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index = index.add %payload_index1, %row_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<4xi32> + %weight_words = vector.load %weight_view[%c0] : view<4xi32> -> vector<4xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %next0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter0, %weight_words, %weight_scale, %group, %c0, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter1, %weight_words, %weight_scale, %group, %c1, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next2 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter2, %weight_words, %weight_scale, %group, %c2, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + scf.yield %next0, %next1, %next2 : f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2 : f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32 + } + + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %partial_view = buffer.view %partial_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %is_partition_zero = index.cmp eq, %partition, %c0 : index + scf.if %valid_row { + scf.if %is_partition_zero { + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + } else { + view.store %result0, %partial_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %partial_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %partial_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + } + } + + kernel.barrier scope(workgroup) ordering(release) + %counter_count = index.div %rows_plus, %c128 : index + %counter_view = buffer.view %counters_na[%base] : buffer -> view<[%counter_count]xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + %one_i32 = scalar.constant 1 : i32 + %zero_i32 = scalar.constant 0 : i32 + %negative_two_i32 = scalar.constant -2 : i32 + %local_old_counter = scf.if %is_arrival_lane -> (i32) { + %old_counter = view.atomic.rmw %one_i32, %counter_view[%workgroup] {ordering = acq_rel, scope = device} : i32, view<[%counter_count]xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %zero_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %one_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + scf.if %valid_row { + %partition_zero0 = view.load %output_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero1 = view.load %output_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero2 = view.load %output_view[%c2, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one0 = view.load %partial_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one1 = view.load %partial_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one2 = view.load %partial_view[%c2, %row] : view<[%cols]x[%rows]xf32> -> f32 + %combined0 = scalar.addf %partition_zero0, %partition_one0 : f32 + %combined1 = scalar.addf %partition_zero1, %partition_one1 : f32 + %combined2 = scalar.addf %partition_zero2, %partition_one2 : f32 + view.store %combined0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_two_i32, %counter_view[%workgroup] {ordering = release, scope = device} : i32, view<[%counter_count]xi32> + } + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c4") @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c4() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %rows_up = index.add %rows, %c127 : index + %row_workgroups = index.div %rows_up, %c128 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %qact: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 4, 4)] : index + %base = index.constant 0 : offset + %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 + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + + %split_workgroup = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %partition = index.rem %split_workgroup, %c2 : index + %workgroup = index.div %split_workgroup, %c2 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %row_base = index.mul %workgroup, %c128 : index + %row = index.add %row_base, %workitem : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %groups_per_partition = index.div %group_count, %c2 : index + %group_begin = index.mul %partition, %groups_per_partition : index + %group_end = index.add %group_begin, %groups_per_partition : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %weight_global = buffer.assume.memory_space %weight : buffer + %output_global = buffer.assume.memory_space %output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %partial_global = buffer.assume.memory_space %partial : buffer + %counters_global = buffer.assume.memory_space %completion_counters : buffer + %weight_na, %output_na, %qact_na, %scales_na, %partial_na, %counters_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global, %partial_global, %counters_global : buffer, buffer, buffer, buffer, buffer, buffer + + %result0, %result1, %result2, %result3 = scf.if %valid_row -> (f32, f32, f32, f32) { + %acc0, %acc1, %acc2, %acc3 = scf.for %group = [%group_begin to %group_end step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32, %iter3 = %zero_f32 : f32) -> (f32, f32, f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index = index.add %payload_index1, %row_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<4xi32> + %weight_words = vector.load %weight_view[%c0] : view<4xi32> -> vector<4xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %next0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter0, %weight_words, %weight_scale, %group, %c0, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter1, %weight_words, %weight_scale, %group, %c1, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next2 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter2, %weight_words, %weight_scale, %group, %c2, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next3 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter3, %weight_words, %weight_scale, %group, %c3, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + scf.yield %next0, %next1, %next2, %next3 : f32, f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2, %acc3 : f32, f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32, f32 + } + + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %partial_view = buffer.view %partial_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %is_partition_zero = index.cmp eq, %partition, %c0 : index + scf.if %valid_row { + scf.if %is_partition_zero { + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result3, %output_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + } else { + view.store %result0, %partial_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %partial_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %partial_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result3, %partial_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + } + } + + kernel.barrier scope(workgroup) ordering(release) + %counter_count = index.div %rows_plus, %c128 : index + %counter_view = buffer.view %counters_na[%base] : buffer -> view<[%counter_count]xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + %one_i32 = scalar.constant 1 : i32 + %zero_i32 = scalar.constant 0 : i32 + %negative_two_i32 = scalar.constant -2 : i32 + %local_old_counter = scf.if %is_arrival_lane -> (i32) { + %old_counter = view.atomic.rmw %one_i32, %counter_view[%workgroup] {ordering = acq_rel, scope = device} : i32, view<[%counter_count]xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %zero_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %one_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + scf.if %valid_row { + %partition_zero0 = view.load %output_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero1 = view.load %output_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero2 = view.load %output_view[%c2, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero3 = view.load %output_view[%c3, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one0 = view.load %partial_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one1 = view.load %partial_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one2 = view.load %partial_view[%c2, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one3 = view.load %partial_view[%c3, %row] : view<[%cols]x[%rows]xf32> -> f32 + %combined0 = scalar.addf %partition_zero0, %partition_one0 : f32 + %combined1 = scalar.addf %partition_zero1, %partition_one1 : f32 + %combined2 = scalar.addf %partition_zero2, %partition_one2 : f32 + %combined3 = scalar.addf %partition_zero3, %partition_one3 : f32 + view.store %combined0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined3, %output_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_two_i32, %counter_view[%workgroup] {ordering = release, scope = device} : i32, view<[%counter_count]xi32> + } + } + kernel.return +} + +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c5") @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_c5() { + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c127 = index.constant 127 : index + %c128 = index.constant 128 : index + %rows_up = index.add %rows, %c127 : index + %row_workgroups = index.div %rows_up, %c128 : index + %workgroups = index.mul %row_workgroups, %c2 : index + kernel.launch.config workgroups(%workgroups, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %qact: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) { + %k0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k = index.assume %k0 [range(%k0, 256, 32768), mul(%k0, 256)] : index + %rows = index.assume %rows0 [range(%rows0, 1, 262144)] : index + %cols = index.assume %cols0 [range(%cols0, 5, 5)] : index + %base = index.constant 0 : offset + %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 + %c128 = index.constant 128 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %zero_f32 = scalar.constant 0.0 : f32 + + %split_workgroup = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %partition = index.rem %split_workgroup, %c2 : index + %workgroup = index.div %split_workgroup, %c2 : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 127)] : index + %row_base = index.mul %workgroup, %c128 : index + %row = index.add %row_base, %workitem : index + %valid_row = index.cmp ult, %row, %rows : index + + %block_count = index.div %k, %c256 : index + %group_count = index.div %k, %c32 : index + %groups_per_partition = index.div %group_count, %c2 : index + %group_begin = index.mul %partition, %groups_per_partition : index + %group_end = index.add %group_begin, %groups_per_partition : index + %row_group0 = config.get @ggml.mul_mat.symmetric_i4.lowrow.row_group_size : index + %rows_plus = index.add %rows, %c255 : index + %row_group = index.assume %row_group0 [range(%row_group0, 32, 262144), mul(%row_group0, 32)] : index + %row_group_id = index.div %row, %row_group : index + %row_lane = index.rem %row, %row_group : index + %cohort_lane = index.div %row_lane, %c4 : index + %payload_plane_bytes = index.mul %row_group, %c16 : index + %payload_block_bytes = index.mul %payload_plane_bytes, %c8 : index + %scale_cohorts = index.div %row_group, %c4 : index + %scale_plane_bytes = index.mul %scale_cohorts, %c16 : index + %block_bytes = index.add %payload_block_bytes, %scale_plane_bytes : index + %row_group_bytes = index.mul %block_count, %block_bytes : index + %group_base = index.mul %row_group_id, %row_group_bytes : index + %scale_region_add = index.mul %block_count, %payload_block_bytes : index + %scale_region = index.add %group_base, %scale_region_add : index + + %weight_global = buffer.assume.memory_space %weight : buffer + %output_global = buffer.assume.memory_space %output : buffer + %qact_global = buffer.assume.memory_space %qact : buffer + %scales_global = buffer.assume.memory_space %scales : buffer + %partial_global = buffer.assume.memory_space %partial : buffer + %counters_global = buffer.assume.memory_space %completion_counters : buffer + %weight_na, %output_na, %qact_na, %scales_na, %partial_na, %counters_na = buffer.assume.noalias %weight_global, %output_global, %qact_global, %scales_global, %partial_global, %counters_global : buffer, buffer, buffer, buffer, buffer, buffer + + %result0, %result1, %result2, %result3, %result4 = scf.if %valid_row -> (f32, f32, f32, f32, f32) { + %acc0, %acc1, %acc2, %acc3, %acc4 = scf.for %group = [%group_begin to %group_end step %c1](%iter0 = %zero_f32 : f32, %iter1 = %zero_f32 : f32, %iter2 = %zero_f32 : f32, %iter3 = %zero_f32 : f32, %iter4 = %zero_f32 : f32) -> (f32, f32, f32, f32, f32) { + %block = index.div %group, %c8 : index + %group_in_block = index.rem %group, %c8 : index + %block_payload_add = index.mul %block, %payload_block_bytes : index + %group_payload_add = index.mul %group_in_block, %payload_plane_bytes : index + %row_payload_add = index.mul %row_lane, %c16 : index + %payload_index0 = index.add %group_base, %block_payload_add : index + %payload_index1 = index.add %payload_index0, %group_payload_add : index + %payload_index = index.add %payload_index1, %row_payload_add : index + %payload_base = index.cast %payload_index : index to offset + %weight_view = buffer.view %weight_na[%payload_base] : buffer -> view<4xi32> + %weight_words = vector.load %weight_view[%c0] : view<4xi32> -> vector<4xi32> + + %block_scale_add = index.mul %block, %scale_plane_bytes : index + %cohort_scale_add = index.mul %cohort_lane, %c16 : index + %group_scale_add = index.mul %group_in_block, %c2 : index + %scale_index0 = index.add %scale_region, %block_scale_add : index + %scale_index1 = index.add %scale_index0, %cohort_scale_add : index + %scale_byte_index = index.add %scale_index1, %group_scale_add : index + %scale_byte_base = index.cast %scale_byte_index : index to offset + %weight_scale_view = buffer.view %weight_na[%scale_byte_base] : buffer -> view<1xf16> + %weight_scale_f16 = view.load %weight_scale_view[%c0] : view<1xf16> -> f16 + %weight_scale = scalar.extf %weight_scale_f16 : f16 to f32 + + %next0 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter0, %weight_words, %weight_scale, %group, %c0, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next1 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter1, %weight_words, %weight_scale, %group, %c1, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next2 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter2, %weight_words, %weight_scale, %group, %c2, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next3 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter3, %weight_words, %weight_scale, %group, %c3, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + %next4 = func.call @ggml_mul_mat_symmetric_i4_lowrow_split_k2_direct_dot_full_dot(%iter4, %weight_words, %weight_scale, %group, %c4, %k, %group_count, %qact_na, %scales_na) : (f32, vector<4xi32>, f32, index, index, index, index, buffer, buffer) -> (f32) + scf.yield %next0, %next1, %next2, %next3, %next4 : f32, f32, f32, f32, f32 + } + scf.yield %acc0, %acc1, %acc2, %acc3, %acc4 : f32, f32, f32, f32, f32 + } else { + scf.yield %zero_f32, %zero_f32, %zero_f32, %zero_f32, %zero_f32 : f32, f32, f32, f32, f32 + } + + %output_view = buffer.view %output_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %partial_view = buffer.view %partial_na[%base] : buffer -> view<[%cols]x[%rows]xf32> + %is_partition_zero = index.cmp eq, %partition, %c0 : index + scf.if %valid_row { + scf.if %is_partition_zero { + view.store %result0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result3, %output_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result4, %output_view[%c4, %row] : f32, view<[%cols]x[%rows]xf32> + } else { + view.store %result0, %partial_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result1, %partial_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result2, %partial_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result3, %partial_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %result4, %partial_view[%c4, %row] : f32, view<[%cols]x[%rows]xf32> + } + } + + kernel.barrier scope(workgroup) ordering(release) + %counter_count = index.div %rows_plus, %c128 : index + %counter_view = buffer.view %counters_na[%base] : buffer -> view<[%counter_count]xi32> + %is_arrival_lane = index.cmp eq, %workitem, %c0 : index + %one_i32 = scalar.constant 1 : i32 + %zero_i32 = scalar.constant 0 : i32 + %negative_two_i32 = scalar.constant -2 : i32 + %local_old_counter = scf.if %is_arrival_lane -> (i32) { + %old_counter = view.atomic.rmw %one_i32, %counter_view[%workgroup] {ordering = acq_rel, scope = device} : i32, view<[%counter_count]xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %zero_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %is_last_partition = scalar.cmpi eq, %old_counter, %one_i32 : i32 + scf.if %is_last_partition { + kernel.barrier scope(workgroup) ordering(acquire) + scf.if %valid_row { + %partition_zero0 = view.load %output_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero1 = view.load %output_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero2 = view.load %output_view[%c2, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero3 = view.load %output_view[%c3, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_zero4 = view.load %output_view[%c4, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one0 = view.load %partial_view[%c0, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one1 = view.load %partial_view[%c1, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one2 = view.load %partial_view[%c2, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one3 = view.load %partial_view[%c3, %row] : view<[%cols]x[%rows]xf32> -> f32 + %partition_one4 = view.load %partial_view[%c4, %row] : view<[%cols]x[%rows]xf32> -> f32 + %combined0 = scalar.addf %partition_zero0, %partition_one0 : f32 + %combined1 = scalar.addf %partition_zero1, %partition_one1 : f32 + %combined2 = scalar.addf %partition_zero2, %partition_one2 : f32 + %combined3 = scalar.addf %partition_zero3, %partition_one3 : f32 + %combined4 = scalar.addf %partition_zero4, %partition_one4 : f32 + view.store %combined0, %output_view[%c0, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined1, %output_view[%c1, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined2, %output_view[%c2, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined3, %output_view[%c3, %row] : f32, view<[%cols]x[%rows]xf32> + view.store %combined4, %output_view[%c4, %row] : f32, view<[%cols]x[%rows]xf32> + } + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_lane { + view.atomic.reduce %negative_two_i32, %counter_view[%workgroup] {ordering = release, scope = device} : i32, view<[%counter_count]xi32> + } + } + kernel.return +} + +// Two K partitions increase independent low-row work while retaining the +// canonical signed-I4 contraction and fixed-order F32 publication path. +kernel.def export("ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma") @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma() { + %unit = index.constant 1 : index + %c2 = index.constant 2 : index + %k_n = index.constant 16 : index + %k_m = index.constant 16 : index + %n_m1 = index.constant 15 : index + %m_m1 = index.constant 15 : index + %wg = index.constant 64 : index + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %rows_up = index.add %rows, %n_m1 : index + %row_tiles = index.div %rows_up, %k_n : index + %row_groups_up = index.add %row_tiles, %unit : index + %row_groups = index.div %row_groups_up, %c2 : index + %split_row_groups = index.mul %row_groups, %c2 : index + %cols_up = index.add %cols, %m_m1 : index + %col_tiles = index.div %cols_up, %k_m : index + kernel.launch.config workgroups(%col_tiles, %split_row_groups, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%weight: buffer, %input: buffer, %output: buffer, %aq: buffer, %as: buffer, %asum: buffer, %partial: buffer, %completion_counters: buffer) { + %base = index.constant 0 : offset + %lds_bytes = index.constant 3072 : offset + %scratch = buffer.alloca align(16) %lds_bytes : buffer + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 40 : index + %k_wstride = index.constant 40 : index + %k_m = index.constant 16 : index + %k_n = index.constant 16 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 1 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %split_k = scalar.constant true : i1 + %compact_q2 = scalar.constant false : i1 + + %k = config.get @ggml.mul_mat.symmetric_i4.lowrow.input_size : index + %rows = config.get @ggml.mul_mat.symmetric_i4.lowrow.output_size : index + %cols = config.get @ggml.mul_mat.symmetric_i4.lowrow.token_count : index + %k_b = index.assume %k [range(%k, 64, 32768)] : index + %rows_b = index.assume %rows [range(%rows, 1, 262144)] : index + %cols_b = index.assume %cols [range(%cols, 1, 32768)] : index + %nchunks = index.div %k_b, %k_kc : index + %kblocks = index.div %k_b, %c32 : index + %partition_chunks = index.div %nchunks, %c2 : index + + %col_tile0 = kernel.workgroup.id : index + %split_row_group0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 511)] : index + %split_row_group = index.assume %split_row_group0 [range(%split_row_group0, 0, 32767)] : index + %row_group = index.div %split_row_group, %c2 : index + %partition = index.rem %split_row_group, %c2 : index + %row_tile0 = index.mul %row_group, %c2 : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 16383)] : index + %chunk_begin = index.mul %partition, %partition_chunks : index + %chunk_end = index.add %chunk_begin, %partition_chunks : index + %tid = index.assume %tid0 [range(%tid0, 0, 63)] : index + %wave = index.div %tid, %k_wave : index + + %weight_g = buffer.assume.memory_space %weight : buffer + %output_g = buffer.assume.memory_space %output : buffer + %aq_g = buffer.assume.memory_space %aq : buffer + %as_g = buffer.assume.memory_space %as : buffer + %asum_g = buffer.assume.memory_space %asum : buffer + %partial_g = buffer.assume.memory_space %partial : buffer + %completion_counters_g = buffer.assume.memory_space %completion_counters : buffer + %weight_na, %output_na, %aq_na, %as_na, %asum_na, %partial_na, %completion_counters_na = buffer.assume.noalias %weight_g, %output_g, %aq_g, %as_g, %asum_g, %partial_g, %completion_counters_g : buffer, buffer, buffer, buffer, buffer, buffer, buffer + template.apply<@ggml.mul_mat.symmetric_i4.lowrow.m16n16.body>(%scratch, %weight_na, %output_na, %aq_na, %as_na, %asum_na, %compact_q2, %split_k, %partial_na, %completion_counters_na, %base, %base, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %k_b, %rows_b, %cols_b, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave, %partition, %chunk_begin, %chunk_end) : (buffer, buffer, buffer, buffer, buffer, buffer, i1, i1, buffer, buffer, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index, index, index, index) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_symmetric_i8_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_symmetric_i8_wmma.loom new file mode 100644 index 000000000000..e755fa01056a --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_symmetric_i8_wmma.loom @@ -0,0 +1,762 @@ +config.decl @ggml.quantize_symmetric_i8_k256.token_count : %value: index where [range(%value, 256, 2048), mul(%value, 256)] + +config.decl @ggml.quantize_symmetric_i8_k256.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] + +func.def inline @ggml_symmetric_i8_wmmai8_stage_values(%token256: i1, %scratch: buffer, %base: offset, %adbase: index, %ad1: index, %values0: vector<32xi8>, %values1: vector<32xi8>) { + scf.if %token256 { + %al_flat = buffer.view %scratch[%base] : buffer -> view<20480xi8> + %ad0_bounded = index.assume %adbase [range(%adbase, 0, 20400)] : index + %ad1_bounded = index.assume %ad1 [range(%ad1, 32, 20432)] : index + vector.store %values0, %al_flat[%ad0_bounded] : vector<32xi8>, view<20480xi8> + vector.store %values1, %al_flat[%ad1_bounded] : vector<32xi8>, view<20480xi8> + } else { + %al_flat = buffer.view %scratch[%base] : buffer -> view<10240xi8> + %ad0_bounded = index.assume %adbase [range(%adbase, 0, 10160)] : index + %ad1_bounded = index.assume %ad1 [range(%ad1, 32, 10192)] : index + vector.store %values0, %al_flat[%ad0_bounded] : vector<32xi8>, view<10240xi8> + vector.store %values1, %al_flat[%ad1_bounded] : vector<32xi8>, view<10240xi8> + } + func.return +} + +func.def inline @ggml_symmetric_i8_wmmai8_stage_metadata(%token256: i1, %scratch: buffer, %as_off: offset, %asum_off: offset, %dst0: index, %dst1: index, %d0: f32, %d1: f32, %sum0: f32, %sum1: f32) { + scf.if %token256 { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<512xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<512xf32> + %dst0_bounded = index.assume %dst0 [range(%dst0, 0, 495)] : index + %dst1_bounded = index.assume %dst1 [range(%dst1, 16, 511)] : index + view.store %d0, %scale_view[%dst0_bounded] : f32, view<512xf32> + view.store %d1, %scale_view[%dst1_bounded] : f32, view<512xf32> + view.store %sum0, %sum_view[%dst0_bounded] : f32, view<512xf32> + view.store %sum1, %sum_view[%dst1_bounded] : f32, view<512xf32> + } else { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<256xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<256xf32> + %dst0_bounded = index.assume %dst0 [range(%dst0, 0, 239)] : index + %dst1_bounded = index.assume %dst1 [range(%dst1, 16, 255)] : index + view.store %d0, %scale_view[%dst0_bounded] : f32, view<256xf32> + view.store %d1, %scale_view[%dst1_bounded] : f32, view<256xf32> + view.store %sum0, %sum_view[%dst0_bounded] : f32, view<256xf32> + view.store %sum1, %sum_view[%dst1_bounded] : f32, view<256xf32> + } + func.return +} + +func.def inline @ggml_symmetric_i8_wmmai8_load_lhs(%token256: i1, %scratch: buffer, %base: offset, %row: index, %column: index, %c16: index) -> (vector<16xi8>) { + %lhs_layout = encoding.layout.strided [80, 1] : encoding + %result = scf.if %token256 -> (vector<16xi8>) { + %al_view = buffer.view %scratch[%base] : buffer -> view<256x64xi8, %lhs_layout> + %values = vector.fragment.load %al_view[%row, %column] shape [%c16, %c16] : view<256x64xi8, %lhs_layout> -> vector<16xi8> + scf.yield %values : vector<16xi8> + } else { + %al_view = buffer.view %scratch[%base] : buffer -> view<128x64xi8, %lhs_layout> + %values = vector.fragment.load %al_view[%row, %column] shape [%c16, %c16] : view<128x64xi8, %lhs_layout> -> vector<16xi8> + scf.yield %values : vector<16xi8> + } + func.return %result : vector<16xi8> +} + +func.def inline @ggml_symmetric_i8_wmmai8_load_metadata(%token256: i1, %scratch: buffer, %as_off: offset, %asum_off: offset, %index: index) -> (vector<8xf32>, vector<8xf32>) { + %scale, %sum = scf.if %token256 -> (vector<8xf32>, vector<8xf32>) { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<512xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<512xf32> + %bounded_index = index.assume %index [range(%index, 0, 504)] : index + %scale_values = vector.load %scale_view[%bounded_index] : view<512xf32> -> vector<8xf32> + %sum_values = vector.load %sum_view[%bounded_index] : view<512xf32> -> vector<8xf32> + scf.yield %scale_values, %sum_values : vector<8xf32>, vector<8xf32> + } else { + %scale_view = buffer.view %scratch[%as_off] : buffer -> view<256xf32> + %sum_view = buffer.view %scratch[%asum_off] : buffer -> view<256xf32> + %bounded_index = index.assume %index [range(%index, 0, 120)] : index + %scale_values = vector.load %scale_view[%bounded_index] : view<256xf32> -> vector<8xf32> + %sum_values = vector.load %sum_view[%bounded_index] : view<256xf32> -> vector<8xf32> + scf.yield %scale_values, %sum_values : vector<8xf32>, vector<8xf32> + } + func.return %scale, %sum : vector<8xf32>, vector<8xf32> +} + +amdgpu.target @ggml_quantize_symmetric_i8_k256_gfx11_wave32 {subgroup_size = 32} + +// Quantize one K256 activation block with two least-squares scale refinements. +kernel.def target(@ggml_quantize_symmetric_i8_k256_gfx11_wave32) export("ggml_quantize_f32_symmetric_i8_k256") @ggml_quantize_f32_symmetric_i8_k256(%token_count: index, %input_size_arg: index) { + %token_capacity = config.get @ggml.quantize_symmetric_i8_k256.token_count : index + %input_size = config.get @ggml.quantize_symmetric_i8_k256.input_size : index + %c256 = index.constant 256 : index + %c32 = index.constant 32 : index + %c1 = index.constant 1 : index + %block_count = index.div %input_size, %c256 : index + %group_count = index.mul %token_capacity, %block_count : index + kernel.launch.config workgroups(%group_count, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index +} launch(%token_count: index, %input_size_arg: index, %input: buffer, %output: buffer) { + %token_capacity0 = config.get @ggml.quantize_symmetric_i8_k256.token_count : index + %input_size0 = config.get @ggml.quantize_symmetric_i8_k256.input_size : index + %token_capacity, %input_size = index.assume %token_capacity0, %input_size0 [range(%token_capacity0, 1, 2048), range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index, index + %bounded_token_count, %bounded_input_size = index.assume %token_count, %input_size_arg [eq(%token_count, %token_capacity), eq(%input_size_arg, %input_size)] : index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c32 = index.constant 32 : index + %c256 = index.constant 256 : index + %c127 = scalar.constant 127.0 : f32 + %cn127 = scalar.constant -127.0 : f32 + %cn1 = scalar.constant -1.0 : f32 + %c1f = scalar.constant 1.0 : f32 + %c0f = scalar.constant 0.0 : f32 + %c0h = scalar.constant 0.0 : f16 + %v127 = vector.splat %c127 : vector<8xf32> + %vn127 = vector.splat %cn127 : vector<8xf32> + %vn1 = vector.splat %cn1 : vector<8xf32> + %base = index.constant 0 : offset + %lane0 = kernel.workitem.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index + %group0 = kernel.workgroup.id : index + %block_count = index.div %bounded_input_size, %c256 : index + %group_count = index.mul %bounded_token_count, %block_count : index + %group = index.assume %group0 [range(%group0, 0, 262143), lt(%group0, %group_count)] : index + %block = index.div %group, %bounded_token_count : index + %token = index.rem %group, %bounded_token_count : index + %token_base = index.mul %token, %bounded_input_size : index + %block_add = index.mul %block, %c256 : index + %lane_add = index.mul %lane, %c8 : index + %input_index0 = index.add %token_base, %block_add : index + %input_index = index.add %input_index0, %lane_add : index + %element_count = index.mul %bounded_token_count, %bounded_input_size : index + %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer + %input_view = buffer.view %input_noalias[%base] : buffer -> view<[%element_count]xf32> + %values = vector.load %input_view[%input_index] : view<[%element_count]xf32> -> vector<8xf32> + %absolute = vector.absf %values : vector<8xf32> + %lane_max = vector.reduce %absolute, %c0f : vector<8xf32>, f32 + %amax = kernel.subgroup.reduce %lane_max : f32 + %nonzero = scalar.cmpf one, %amax, %c0f : f32 + %scale0 = scf.if %nonzero -> (f32) { + %value = scalar.divf %amax, %c127 : f32 + scf.yield %value : f32 + } else { + scf.yield %c0f : f32 + } + %inv0 = scf.if %nonzero -> (f32) { + %value = scalar.divf %c1f, %scale0 : f32 + scf.yield %value : f32 + } else { + scf.yield %c0f : f32 + } + %inv0v = vector.splat %inv0 : vector<8xf32> + %scaled0 = vector.mulf %values, %inv0v : vector<8xf32> + %rounded0 = vector.roundf %scaled0 : vector<8xf32> + %lower0 = vector.maxnumf %rounded0, %vn127 : vector<8xf32> + %neg0 = vector.mulf %lower0, %vn1 : vector<8xf32> + %negclamp0 = vector.maxnumf %neg0, %vn127 : vector<8xf32> + %clamped0 = vector.mulf %negclamp0, %vn1 : vector<8xf32> + %numerator_terms0 = vector.mulf %values, %clamped0 : vector<8xf32> + %denominator_terms0 = vector.mulf %clamped0, %clamped0 : vector<8xf32> + %lane_numerator0 = vector.reduce %numerator_terms0, %c0f : vector<8xf32>, f32 + %lane_denominator0 = vector.reduce %denominator_terms0, %c0f : vector<8xf32>, f32 + %numerator0 = kernel.subgroup.reduce %lane_numerator0 : f32 + %denominator0 = kernel.subgroup.reduce %lane_denominator0 : f32 + %has_denominator0 = scalar.cmpf one, %denominator0, %c0f : f32 + %scale1 = scf.if %has_denominator0 -> (f32) { + %value = scalar.divf %numerator0, %denominator0 : f32 + scf.yield %value : f32 + } else { + scf.yield %scale0 : f32 + } + %inv1 = scf.if %nonzero -> (f32) { + %value = scalar.divf %c1f, %scale1 : f32 + scf.yield %value : f32 + } else { + scf.yield %c0f : f32 + } + %inv1v = vector.splat %inv1 : vector<8xf32> + %scaled1 = vector.mulf %values, %inv1v : vector<8xf32> + %rounded1 = vector.roundf %scaled1 : vector<8xf32> + %lower1 = vector.maxnumf %rounded1, %vn127 : vector<8xf32> + %neg1 = vector.mulf %lower1, %vn1 : vector<8xf32> + %negclamp1 = vector.maxnumf %neg1, %vn127 : vector<8xf32> + %clamped1 = vector.mulf %negclamp1, %vn1 : vector<8xf32> + %numerator_terms1 = vector.mulf %values, %clamped1 : vector<8xf32> + %denominator_terms1 = vector.mulf %clamped1, %clamped1 : vector<8xf32> + %lane_numerator1 = vector.reduce %numerator_terms1, %c0f : vector<8xf32>, f32 + %lane_denominator1 = vector.reduce %denominator_terms1, %c0f : vector<8xf32>, f32 + %numerator1 = kernel.subgroup.reduce %lane_numerator1 : f32 + %denominator1 = kernel.subgroup.reduce %lane_denominator1 : f32 + %has_denominator1 = scalar.cmpf one, %denominator1, %c0f : f32 + %scale2 = scf.if %has_denominator1 -> (f32) { + %value = scalar.divf %numerator1, %denominator1 : f32 + scf.yield %value : f32 + } else { + scf.yield %scale1 : f32 + } + %encoded_scale = scalar.fptrunc %scale2 : f32 to f16 + %decoded_scale = scalar.extf %encoded_scale : f16 to f32 + %inverse = scf.if %nonzero -> (f32) { + %value = scalar.divf %c1f, %decoded_scale : f32 + scf.yield %value : f32 + } else { + scf.yield %c0f : f32 + } + %inverse_v = vector.splat %inverse : vector<8xf32> + %scaled = vector.mulf %values, %inverse_v : vector<8xf32> + %rounded = vector.roundf %scaled : vector<8xf32> + %lower = vector.maxnumf %rounded, %vn127 : vector<8xf32> + %neg = vector.mulf %lower, %vn1 : vector<8xf32> + %negclamp = vector.maxnumf %neg, %vn127 : vector<8xf32> + %clamped = vector.mulf %negclamp, %vn1 : vector<8xf32> + %quantized = vector.fptosi %clamped : vector<8xf32> to vector<8xi8> + %chunk = index.div %lane, %c4 : index + %word = index.rem %lane, %c4 : index + %block_chunk0 = index.mul %block, %c8 : index + %block_chunk = index.add %block_chunk0, %chunk : index + %payload_row0 = index.mul %block_chunk, %bounded_token_count : index + %payload_row = index.add %payload_row0, %token : index + %payload_base = index.mul %payload_row, %c32 : index + %word_add = index.mul %word, %c8 : index + %payload_index = index.add %payload_base, %word_add : index + %output_payload = buffer.view %output_noalias[%base] : buffer -> view<[%element_count]xi8> + vector.store %quantized, %output_payload[%payload_index] : vector<8xi8>, view<[%element_count]xi8> + %metadata_offset = index.cast %element_count : index to offset + %metadata = buffer.view %output_noalias[%metadata_offset] : buffer -> view<[%group_count]xi32> + %is_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_leader { + %metadata_pair = vector.from_elements %encoded_scale, %c0h : vector<2xf16> + %metadata_word = vector.bitcast %metadata_pair : vector<2xf16> to vector<1xi32> + vector.store %metadata_word, %metadata[%group] : vector<1xi32>, view<[%group_count]xi32> + } + kernel.return +} + +kernel.def target(@ggml_symmetric_i8_gfx11_wave32) export("ggml_mul_mat_q5_k_symmetric_i8_wmma") @ggml_mul_mat_q5_k_symmetric_i8_wmma(%token_count: index) { + %unit = index.constant 1 : index + %tile_rows = index.constant 64 : index + %tile_tokens = index.constant 256 : index + %workgroup_size = index.constant 256 : index + %output_size = config.get @ggml.mul_mat.symmetric_i8.output_size : index + %token_capacity = config.get @ggml.mul_mat.symmetric_i8.token_count : index + %row_tiles = index.div %output_size, %tile_rows : index + %token_tiles = index.div %token_capacity, %tile_tokens : index + kernel.launch.config workgroups(%token_tiles, %row_tiles, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) { + %base = index.constant 0 : offset + %lds_bytes = index.constant 30720 : offset + %w_off = index.constant 20480 : offset + %as_off = index.constant 25600 : offset + %ws_off = index.constant 27648 : offset + %asum_off = index.constant 28160 : offset + %wc_off = index.constant 30208 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c17 = index.constant 17 : index + %c34 = index.constant 34 : index + %k_kc = index.constant 64 : index + %k_astride = index.constant 80 : index + %k_wstride = index.constant 80 : index + %k_m = index.constant 256 : index + %k_n = index.constant 64 : index + %k_blocks = index.constant 2 : index + %k_wave = index.constant 32 : index + %k_wm = index.constant 8 : index + %fzero = vector.constant 0.0 : vector<8xf32> + %izero = vector.constant 0 : vector<8xi32> + %is_iq4xs = scalar.constant false : i1 + %token256 = scalar.constant true : i1 + %q8_plane = scalar.constant true : i1 + %apply_swiglu = scalar.constant false : i1 + + %input_size0 = config.get @ggml.mul_mat.symmetric_i8.input_size : index + %output_size0 = config.get @ggml.mul_mat.symmetric_i8.output_size : index + %token_capacity = config.get @ggml.mul_mat.symmetric_i8.token_count : index + %input_size = index.assume %input_size0 [range(%input_size0, 256, 32768), mul(%input_size0, 256)] : index + %output_size = index.assume %output_size0 [range(%output_size0, 64, 262144), mul(%output_size0, 64)] : index + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %token_capacity [range(%token_count, 256, 2048), mul(%token_count, 256), eq(%token_count, %token_capacity)] : index, index + %nchunks = index.div %input_size, %k_kc : index + %kblocks = index.div %input_size, %c32 : index + + %col_tile0 = kernel.workgroup.id : index + %row_tile0 = kernel.workgroup.id : index + %tid0 = kernel.workitem.id : index + %col_tile = index.assume %col_tile0 [range(%col_tile0, 0, 7)] : index + %row_tile = index.assume %row_tile0 [range(%row_tile0, 0, 4095)] : index + %tid = index.assume %tid0 [range(%tid0, 0, 255)] : index + %wave = index.div %tid, %k_wave : index + + %weight_g = buffer.assume.memory_space %weight : buffer + %output_g = buffer.assume.memory_space %output : buffer + %q8_g = buffer.assume.memory_space %q8_input : buffer + %weight_noalias, %output_noalias, %q8_noalias = buffer.assume.noalias %weight_g, %output_g, %q8_g : buffer, buffer, buffer + + template.apply<@ggml.mul_mat.symmetric_i8.wmmai8.body>(%is_iq4xs, %token256, %q8_plane, %apply_swiglu, %output, %weight_noalias, %output_noalias, %q8_noalias, %base, %lds_bytes, %w_off, %as_off, %ws_off, %asum_off, %wc_off, %c0, %c1, %c2, %c4, %c8, %c16, %c32, %c17, %c34, %k_kc, %k_astride, %k_wstride, %k_m, %k_n, %k_blocks, %k_wave, %k_wm, %fzero, %izero, %input_size, %output_size, %bounded_token_count, %nchunks, %kblocks, %col_tile, %row_tile, %tid, %wave) : (i1, i1, i1, i1, buffer, buffer, buffer, buffer, offset, offset, offset, offset, offset, offset, offset, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, index, vector<8xf32>, vector<8xi32>, index, index, index, index, index, index, index, index, index) + kernel.return +} + +template.decl @ggml.mul_mat.symmetric_i8.wmmai8.body(%is_iq4xs: i1, %token256: i1, %q8_plane: i1, %apply_swiglu: i1, %gate: buffer, %src0_na: buffer, %dst_na: buffer, %q8_na: buffer, %base: offset, %lds_bytes: offset, %w_off: offset, %as_off: offset, %ws_off: offset, %asum_off: offset, %wc_off: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index) + +amdgpu.target @ggml_symmetric_i8_gfx11_wave32 {subgroup_size = 32} + +config.def @ggml.mul_mat.symmetric_i8.input_size = 5120 : index + +config.def @ggml.mul_mat.symmetric_i8.output_size = 17408 : index + +config.def @ggml.mul_mat.symmetric_i8.token_count = 512 : index + +config.def @ggml.mul_mat.symmetric_i8.activation_cadence = 256 : index + +config.def @ggml.mul_mat.symmetric_i8.chunks_per_activation = 4 : index + +config.def @ggml.mul_mat.symmetric_i8.activation_blocks_per_weight = 1 : index + +// Raw Q5_K x packed-Q8_1_x4 IU8 WMMA using the existing Loom activation layout. +template.def<@ggml.mul_mat.symmetric_i8.wmmai8.body> device @ggml_mul_mat_symmetric_i8_wmmai8_body(%is_iq4xs: i1, %token256: i1, %q8_plane: i1, %apply_swiglu: i1, %gate: buffer, %src0_na: buffer, %dst_na: buffer, %q8_na: buffer, %base: offset, %lds_bytes: offset, %w_off: offset, %as_off: offset, %ws_off: offset, %asum_off: offset, %wc_off: offset, %c0: index, %c1: index, %c2: index, %c4: index, %c8: index, %c16: index, %c32: index, %c17: index, %c34: index, %k_kc: index, %k_astride: index, %k_wstride: index, %k_m: index, %k_n: index, %k_blocks: index, %k_wave: index, %k_wm: index, %fzero: vector<8xf32>, %izero: vector<8xi32>, %k_b: index, %rows_b: index, %cols_b: index, %nchunks: index, %kblocks: index, %col_tile: index, %row_tile: index, %tid: index, %wave: index) { + %q4_k_block = index.constant 256 : index + %q5_block_bytes = index.constant 176 : offset + %q5_high_offset = index.constant 16 : offset + %q5_code_offset = index.constant 48 : offset + %q8_group_bytes = index.constant 144 : offset + %q8_payload_offset = index.constant 16 : offset + %q8_elements_per_group = index.constant 128 : index + // dst is [rows, cols] with rows contiguous, and C is [cols, rows]. + %dst_view = buffer.view %dst_na[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + %gate_global = buffer.assume.memory_space %gate : buffer + %gate_view = buffer.view %gate_global[%base] : buffer -> view<[%cols_b]x[%rows_b]xf32> + + // Four or eight waves share one 64-row weight tile across the token tile. + %scratch = buffer.alloca align(16) %lds_bytes : buffer + %i8_schema = encoding.define #encoding.operand : encoding + // A compile-time format selects unsigned Q5_K or signed IQ4_XS operands before target lowering. + %signed_rhs_schema = encoding.define #encoding.operand : encoding + %unsigned_rhs_schema = encoding.define #encoding.operand : encoding + %u8_schema = scf.select %is_iq4xs, %signed_rhs_schema, %unsigned_rhs_schema : encoding + %rhs_layout = encoding.layout.strided [1, 80] : encoding + %wl_view = buffer.view %scratch[%w_off] : buffer -> view<64x64xi8, %rhs_layout> + %wl_flat = buffer.view %scratch[%w_off] : buffer -> view<5120xi8> + %wsl_view = buffer.view %scratch[%ws_off] : buffer -> view<128xf32> + %wcl_view = buffer.view %scratch[%wc_off] : buffer -> view<128xf32> + + %col_base = index.mul %col_tile, %k_m : index + %row_base = index.mul %row_tile, %k_n : index + + // One thread stages both K32 blocks for a token from the 144-byte Q8_1_x4 group. + %ascol_g = index.add %col_base, %tid : index + %adbase = index.mul %tid, %k_astride : index + %q8_groups_per_row = index.div %k_b, %q8_elements_per_group : index + %q8_row_bytes = index.scale %q8_groups_per_row, %q8_group_bytes : index, offset -> offset + %q8_row_byte_base = index.scale %ascol_g, %q8_row_bytes : index, offset -> offset + %q8_plane_payload_elements = index.mul %k_b, %cols_b : index + %q8_plane_payload_words = index.div %q8_plane_payload_elements, %c4 : index + %q8_plane_metadata_offset = index.cast %q8_plane_payload_elements : index to offset + %q8_plane_metadata_count = index.mul %kblocks, %cols_b : index + %q8_plane_payload = buffer.view %q8_na[%base] : buffer -> view<[%q8_plane_payload_words]xi32> + %q8_plane_metadata = buffer.view %q8_na[%q8_plane_metadata_offset] : buffer -> view<[%q8_plane_metadata_count]xi32> + + // The waves own 32-token blocks and share one 64-row weight tile. + %wcol = index.rem %wave, %k_wm : index + %wrow = index.div %wave, %k_wm : index + %k_mspan = index.constant 32 : index + %k_nspan = index.constant 64 : index + %wm_off = index.mul %wcol, %k_mspan : index + %wn_off = index.mul %wrow, %k_nspan : index + %m_out = index.add %col_base, %wm_off : index + %n_out = index.add %row_base, %wn_off : index + %lm0 = index.add %wm_off, %c0 : index + %gm0 = index.add %m_out, %c0 : index + %k_ma1 = index.constant 16 : index + %lm1 = index.add %wm_off, %k_ma1 : index + %gm1 = index.add %m_out, %k_ma1 : index + %ln0 = index.add %wn_off, %c0 : index + %gn0 = index.add %n_out, %c0 : index + %k_nb1 = index.constant 16 : index + %ln1 = index.add %wn_off, %k_nb1 : index + %gn1 = index.add %n_out, %k_nb1 : index + %k_nb2 = index.constant 32 : index + %ln2 = index.add %wn_off, %k_nb2 : index + %gn2 = index.add %n_out, %k_nb2 : index + %k_nb3 = index.constant 48 : index + %ln3 = index.add %wn_off, %k_nb3 : index + %gn3 = index.add %n_out, %k_nb3 : index + %lane = index.rem %tid, %k_wave : index + %lane_lo = index.rem %lane, %c16 : index + %lane_hi = index.div %lane, %c16 : index + %wsb0_0 = index.add %ln0, %lane_lo : index + %wsb0 = index.mul %wsb0_0, %k_blocks : index + %wsb1_0 = index.add %ln1, %lane_lo : index + %wsb1 = index.mul %wsb1_0, %k_blocks : index + %wsb2_0 = index.add %ln2, %lane_lo : index + %wsb2 = index.mul %wsb2_0, %k_blocks : index + %wsb3_0 = index.add %ln3, %lane_lo : index + %wsb3 = index.mul %wsb3_0, %k_blocks : index + // Activation scale metadata is [tile][block][lane-half][register]. + %as_tile0 = index.div %lm0, %c16 : index + %asb0 = index.mul %as_tile0, %k_blocks : index + %as_tile1 = index.div %lm1, %c16 : index + %asb1 = index.mul %as_tile1, %k_blocks : index + + %i8_c64 = index.constant 64 : index + %i8_c258 = index.constant 258 : index + %i8_c256 = index.constant 256 : index + %i8_act_cadence = config.get @ggml.mul_mat.symmetric_i8.activation_cadence : index + %i8_chunks_per_act = config.get @ggml.mul_mat.symmetric_i8.chunks_per_activation : index + %i8_act_blocks_per_weight = config.get @ggml.mul_mat.symmetric_i8.activation_blocks_per_weight : index + %i8_c48_outer = index.constant 48 : index + %i8_record_bytes = index.constant 16512 : index + %i8_scale_bytes = index.constant 128 : index + %i8_field_bytes = index.constant 4096 : index + %i8_payload_elements = index.mul %k_b, %cols_b : index + %i8_payload_words = index.div %i8_payload_elements, %c4 : index + %i8_meta_offset = index.cast %i8_payload_elements : index to offset + %i8_act_block_count = index.div %k_b, %i8_act_cadence : index + %i8_weight_block_count = index.div %k_b, %i8_c256 : index + %i8_meta_count = index.mul %i8_act_block_count, %cols_b : index + %i8_q8_payload = buffer.view %q8_na[%base] : buffer -> view<[%i8_payload_words]xi32> + %i8_q8_meta = buffer.view %q8_na[%i8_meta_offset] : buffer -> view<[%i8_meta_count]xi32> + %i8_weight_records = index.mul %rows_b, %i8_weight_block_count : index + %i8_weight_bytes = index.mul %i8_weight_records, %i8_c258 : index + %i8_weight_view = buffer.view %src0_na[%base] : buffer -> view<[%i8_weight_bytes]xi8> + %f0_0, %f0_1, %f0_2, %f0_3, %f1_0, %f1_1, %f1_2, %f1_3 = scf.for %i8_act_block = [%c0 to %i8_act_block_count step %c1](%fc0_0 = %fzero : vector<8xf32>, %fc0_1 = %fzero : vector<8xf32>, %fc0_2 = %fzero : vector<8xf32>, %fc0_3 = %fzero : vector<8xf32>, %fc1_0 = %fzero : vector<8xf32>, %fc1_1 = %fzero : vector<8xf32>, %fc1_2 = %fzero : vector<8xf32>, %fc1_3 = %fzero : vector<8xf32>) -> (vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>) { + %i8_source_row_base = index.constant 0 : index + %i8_weight_block = index.div %i8_act_block, %i8_act_blocks_per_weight : index + %i8_act_in_weight = index.rem %i8_act_block, %i8_act_blocks_per_weight : index + %i8_weight_inner_base = index.mul %i8_act_in_weight, %i8_chunks_per_act : index + %i8_tile_block0 = index.mul %row_tile, %i8_weight_block_count : index + %i8_tile_block = index.add %i8_tile_block0, %i8_weight_block : index + %i8_record_base = index.mul %i8_tile_block, %i8_record_bytes : index + %i8_i0_0, %i8_i0_1, %i8_i0_2, %i8_i0_3, %i8_i1_0, %i8_i1_1, %i8_i1_2, %i8_i1_3 = scf.for %inner = [%c0 to %i8_chunks_per_act step %c1](%i8_ic0_0 = %izero : vector<8xi32>, %i8_ic0_1 = %izero : vector<8xi32>, %i8_ic0_2 = %izero : vector<8xi32>, %i8_ic0_3 = %izero : vector<8xi32>, %i8_ic1_0 = %izero : vector<8xi32>, %i8_ic1_1 = %izero : vector<8xi32>, %i8_ic1_2 = %izero : vector<8xi32>, %i8_ic1_3 = %izero : vector<8xi32>) -> (vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>) { + %i8_chunk0_loop = index.mul %i8_act_block, %i8_chunks_per_act : index + %i8_chunk_loop = index.add %i8_chunk0_loop, %inner : index + %i8_q8_block0_loop = index.mul %i8_chunk_loop, %c2 : index + %i8_q8_block1_loop = index.add %i8_q8_block0_loop, %c1 : index + %i8_q8_token_base0_loop = index.mul %i8_q8_block0_loop, %cols_b : index + %i8_q8_token_base1_loop = index.mul %i8_q8_block1_loop, %cols_b : index + %i8_q8_token0_loop = index.add %i8_q8_token_base0_loop, %ascol_g : index + %i8_q8_token1_loop = index.add %i8_q8_token_base1_loop, %ascol_g : index + %i8_q8_word00_loop = index.mul %i8_q8_token0_loop, %c8 : index + %i8_q8_word10_loop = index.mul %i8_q8_token1_loop, %c8 : index + %i8_q8_word0_loop = index.assume %i8_q8_word00_loop [range(%i8_q8_word00_loop, 0, 268435448), mul(%i8_q8_word00_loop, 8)] : index + %i8_q8_word1_loop = index.assume %i8_q8_word10_loop [range(%i8_q8_word10_loop, 8, 268435448), mul(%i8_q8_word10_loop, 8)] : index + %i8_q8_words0_loop = vector.load %i8_q8_payload[%i8_q8_word0_loop] : view<[%i8_payload_words]xi32> -> vector<8xi32> + %i8_q8_words1_loop = vector.load %i8_q8_payload[%i8_q8_word1_loop] : view<[%i8_payload_words]xi32> -> vector<8xi32> + %i8_q8_values0_loop = vector.bitcast %i8_q8_words0_loop : vector<8xi32> to vector<32xi8> + %i8_q8_values1_loop = vector.bitcast %i8_q8_words1_loop : vector<8xi32> to vector<32xi8> + %i8_ad1_loop = index.add %adbase, %c32 : index + func.call @ggml_symmetric_i8_wmmai8_stage_values(%token256, %scratch, %base, %adbase, %i8_ad1_loop, %i8_q8_values0_loop, %i8_q8_values1_loop) : (i1, buffer, offset, index, index, vector<32xi8>, vector<32xi8>) + %i8_meta_block_base_loop = index.mul %i8_act_block, %cols_b : index + %i8_meta_index_loop = index.add %i8_meta_block_base_loop, %ascol_g : index + %i8_meta_word_loop = view.load %i8_q8_meta[%i8_meta_index_loop] : view<[%i8_meta_count]xi32> -> i32 + %i8_meta_vector_loop = vector.from_elements %i8_meta_word_loop : vector<1xi32> + %i8_meta_halves_loop = vector.bitcast %i8_meta_vector_loop : vector<1xi32> to vector<2xf16> + %i8_qscale_f16_loop = vector.extract %i8_meta_halves_loop[0] : vector<2xf16> -> f16 + %i8_qscale_loop = scalar.extf %i8_qscale_f16_loop : f16 to f32 + %i8_qtoken_tile_loop = index.div %tid, %c16 : index + %i8_qtoken_inner_loop = index.rem %tid, %c16 : index + %i8_qtoken_parity_loop = index.rem %i8_qtoken_inner_loop, %c2 : index + %i8_qtoken_vector_loop = index.div %i8_qtoken_inner_loop, %c2 : index + %i8_qscale_tile0_loop = index.mul %i8_qtoken_tile_loop, %c4 : index + %i8_qscale_tile1_loop = index.add %i8_qscale_tile0_loop, %i8_qtoken_parity_loop : index + %i8_qscale_tile2_loop = index.mul %i8_qscale_tile1_loop, %c8 : index + %i8_qscale_dst0_loop = index.add %i8_qscale_tile2_loop, %i8_qtoken_vector_loop : index + %i8_qscale_dst1_loop = index.add %i8_qscale_dst0_loop, %c16 : index + func.call @ggml_symmetric_i8_wmmai8_stage_metadata(%token256, %scratch, %as_off, %asum_off, %i8_qscale_dst0_loop, %i8_qscale_dst1_loop, %i8_qscale_loop, %i8_qscale_loop, %i8_qscale_loop, %i8_qscale_loop) : (i1, buffer, offset, offset, index, index, f32, f32, f32, f32) + %i8_word0 = index.rem %tid, %c4 : index + %i8_word = index.assume %i8_word0 [range(%i8_word0, 0, 3)] : index + %i8_byte_position0 = index.mul %i8_word, %c16 : index + %i8_byte_position = index.assume %i8_byte_position0 [range(%i8_byte_position0, 0, 48), mul(%i8_byte_position0, 16)] : index + %i8_local_row0 = index.div %tid, %c4 : index + %i8_local_row = index.assume %i8_local_row0 [range(%i8_local_row0, 0, 63)] : index + %i8_weight_inner = index.add %i8_weight_inner_base, %inner : index + %i8_inner_bytes = index.mul %i8_weight_inner, %i8_field_bytes : index + %i8_row_byte0 = index.mul %i8_local_row, %i8_c64 : index + %i8_row_byte = index.add %i8_row_byte0, %i8_byte_position : index + %i8_payload_base = index.add %i8_record_base, %i8_scale_bytes : index + %i8_chunk_base = index.add %i8_payload_base, %i8_inner_bytes : index + %i8_address = index.add %i8_chunk_base, %i8_row_byte : index + %i8_values = vector.load %i8_weight_view[%i8_address] : view<[%i8_weight_bytes]xi8> -> vector<16xi8> + %i8_weight_row_base0 = index.mul %i8_local_row, %k_wstride : index + %i8_weight_row_base = index.add %i8_weight_row_base0, %i8_byte_position : index + %i8_weight_destination = index.assume %i8_weight_row_base [range(%i8_weight_row_base, 0, 5104), mul(%i8_weight_row_base, 16)] : index + vector.store %i8_values, %wl_flat[%i8_weight_destination] : vector<16xi8>, view<5120xi8> + kernel.barrier scope(workgroup) ordering(acq_rel) + %i8_lhs0_0_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %c0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs0_0_words = vector.bitcast %i8_lhs0_0_raw : vector<16xi8> to vector<4xi32> + %i8_lhs0_0 = vector.fragment %i8_lhs0_0_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs0_1_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %c16, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs0_1_words = vector.bitcast %i8_lhs0_1_raw : vector<16xi8> to vector<4xi32> + %i8_lhs0_1 = vector.fragment %i8_lhs0_1_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs0_2_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %c32, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs0_2_words = vector.bitcast %i8_lhs0_2_raw : vector<16xi8> to vector<4xi32> + %i8_lhs0_2 = vector.fragment %i8_lhs0_2_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs0_3_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm0, %i8_c48_outer, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs0_3_words = vector.bitcast %i8_lhs0_3_raw : vector<16xi8> to vector<4xi32> + %i8_lhs0_3 = vector.fragment %i8_lhs0_3_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs1_0_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %c0, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs1_0_words = vector.bitcast %i8_lhs1_0_raw : vector<16xi8> to vector<4xi32> + %i8_lhs1_0 = vector.fragment %i8_lhs1_0_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs1_1_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %c16, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs1_1_words = vector.bitcast %i8_lhs1_1_raw : vector<16xi8> to vector<4xi32> + %i8_lhs1_1 = vector.fragment %i8_lhs1_1_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs1_2_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %c32, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs1_2_words = vector.bitcast %i8_lhs1_2_raw : vector<16xi8> to vector<4xi32> + %i8_lhs1_2 = vector.fragment %i8_lhs1_2_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_lhs1_3_raw = func.call @ggml_symmetric_i8_wmmai8_load_lhs(%token256, %scratch, %base, %lm1, %i8_c48_outer, %c16) : (i1, buffer, offset, index, index, index) -> (vector<16xi8>) + %i8_lhs1_3_words = vector.bitcast %i8_lhs1_3_raw : vector<16xi8> to vector<4xi32> + %i8_lhs1_3 = vector.fragment %i8_lhs1_3_words shape [%c16, %c16] using {schema = %i8_schema : encoding} : vector<4xi32> + %i8_rhs0_0_raw = vector.fragment.load %wl_view[%c0, %ln0] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs0_0_words = vector.bitcast %i8_rhs0_0_raw : vector<16xi8> to vector<4xi32> + %i8_rhs0_0 = vector.fragment %i8_rhs0_0_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs0_1_raw = vector.fragment.load %wl_view[%c16, %ln0] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs0_1_words = vector.bitcast %i8_rhs0_1_raw : vector<16xi8> to vector<4xi32> + %i8_rhs0_1 = vector.fragment %i8_rhs0_1_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs0_2_raw = vector.fragment.load %wl_view[%c32, %ln0] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs0_2_words = vector.bitcast %i8_rhs0_2_raw : vector<16xi8> to vector<4xi32> + %i8_rhs0_2 = vector.fragment %i8_rhs0_2_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs0_3_raw = vector.fragment.load %wl_view[%i8_c48_outer, %ln0] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs0_3_words = vector.bitcast %i8_rhs0_3_raw : vector<16xi8> to vector<4xi32> + %i8_rhs0_3 = vector.fragment %i8_rhs0_3_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs1_0_raw = vector.fragment.load %wl_view[%c0, %ln1] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs1_0_words = vector.bitcast %i8_rhs1_0_raw : vector<16xi8> to vector<4xi32> + %i8_rhs1_0 = vector.fragment %i8_rhs1_0_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs1_1_raw = vector.fragment.load %wl_view[%c16, %ln1] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs1_1_words = vector.bitcast %i8_rhs1_1_raw : vector<16xi8> to vector<4xi32> + %i8_rhs1_1 = vector.fragment %i8_rhs1_1_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs1_2_raw = vector.fragment.load %wl_view[%c32, %ln1] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs1_2_words = vector.bitcast %i8_rhs1_2_raw : vector<16xi8> to vector<4xi32> + %i8_rhs1_2 = vector.fragment %i8_rhs1_2_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs1_3_raw = vector.fragment.load %wl_view[%i8_c48_outer, %ln1] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs1_3_words = vector.bitcast %i8_rhs1_3_raw : vector<16xi8> to vector<4xi32> + %i8_rhs1_3 = vector.fragment %i8_rhs1_3_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs2_0_raw = vector.fragment.load %wl_view[%c0, %ln2] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs2_0_words = vector.bitcast %i8_rhs2_0_raw : vector<16xi8> to vector<4xi32> + %i8_rhs2_0 = vector.fragment %i8_rhs2_0_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs2_1_raw = vector.fragment.load %wl_view[%c16, %ln2] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs2_1_words = vector.bitcast %i8_rhs2_1_raw : vector<16xi8> to vector<4xi32> + %i8_rhs2_1 = vector.fragment %i8_rhs2_1_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs2_2_raw = vector.fragment.load %wl_view[%c32, %ln2] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs2_2_words = vector.bitcast %i8_rhs2_2_raw : vector<16xi8> to vector<4xi32> + %i8_rhs2_2 = vector.fragment %i8_rhs2_2_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs2_3_raw = vector.fragment.load %wl_view[%i8_c48_outer, %ln2] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs2_3_words = vector.bitcast %i8_rhs2_3_raw : vector<16xi8> to vector<4xi32> + %i8_rhs2_3 = vector.fragment %i8_rhs2_3_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs3_0_raw = vector.fragment.load %wl_view[%c0, %ln3] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs3_0_words = vector.bitcast %i8_rhs3_0_raw : vector<16xi8> to vector<4xi32> + %i8_rhs3_0 = vector.fragment %i8_rhs3_0_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs3_1_raw = vector.fragment.load %wl_view[%c16, %ln3] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs3_1_words = vector.bitcast %i8_rhs3_1_raw : vector<16xi8> to vector<4xi32> + %i8_rhs3_1 = vector.fragment %i8_rhs3_1_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs3_2_raw = vector.fragment.load %wl_view[%c32, %ln3] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs3_2_words = vector.bitcast %i8_rhs3_2_raw : vector<16xi8> to vector<4xi32> + %i8_rhs3_2 = vector.fragment %i8_rhs3_2_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_rhs3_3_raw = vector.fragment.load %wl_view[%i8_c48_outer, %ln3] shape [%c16, %c16] : view<64x64xi8, %rhs_layout> -> vector<16xi8> + %i8_rhs3_3_words = vector.bitcast %i8_rhs3_3_raw : vector<16xi8> to vector<4xi32> + %i8_rhs3_3 = vector.fragment %i8_rhs3_3_words shape [%c16, %c16] using {schema = %signed_rhs_schema : encoding} : vector<4xi32> + %i8_init0_0 = vector.fragment %i8_ic0_0 shape [%c16, %c16] : vector<8xi32> + %i8_mma0_0_0 = vector.mma %i8_lhs0_0, %i8_rhs0_0, %i8_init0_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_0_1 = vector.mma %i8_lhs0_1, %i8_rhs0_1, %i8_mma0_0_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_0_2 = vector.mma %i8_lhs0_2, %i8_rhs0_2, %i8_mma0_0_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_0_3 = vector.mma %i8_lhs0_3, %i8_rhs0_3, %i8_mma0_0_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init0_1 = vector.fragment %i8_ic0_1 shape [%c16, %c16] : vector<8xi32> + %i8_mma0_1_0 = vector.mma %i8_lhs0_0, %i8_rhs1_0, %i8_init0_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_1_1 = vector.mma %i8_lhs0_1, %i8_rhs1_1, %i8_mma0_1_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_1_2 = vector.mma %i8_lhs0_2, %i8_rhs1_2, %i8_mma0_1_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_1_3 = vector.mma %i8_lhs0_3, %i8_rhs1_3, %i8_mma0_1_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init0_2 = vector.fragment %i8_ic0_2 shape [%c16, %c16] : vector<8xi32> + %i8_mma0_2_0 = vector.mma %i8_lhs0_0, %i8_rhs2_0, %i8_init0_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_2_1 = vector.mma %i8_lhs0_1, %i8_rhs2_1, %i8_mma0_2_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_2_2 = vector.mma %i8_lhs0_2, %i8_rhs2_2, %i8_mma0_2_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_2_3 = vector.mma %i8_lhs0_3, %i8_rhs2_3, %i8_mma0_2_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init0_3 = vector.fragment %i8_ic0_3 shape [%c16, %c16] : vector<8xi32> + %i8_mma0_3_0 = vector.mma %i8_lhs0_0, %i8_rhs3_0, %i8_init0_3 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_3_1 = vector.mma %i8_lhs0_1, %i8_rhs3_1, %i8_mma0_3_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_3_2 = vector.mma %i8_lhs0_2, %i8_rhs3_2, %i8_mma0_3_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma0_3_3 = vector.mma %i8_lhs0_3, %i8_rhs3_3, %i8_mma0_3_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init1_0 = vector.fragment %i8_ic1_0 shape [%c16, %c16] : vector<8xi32> + %i8_mma1_0_0 = vector.mma %i8_lhs1_0, %i8_rhs0_0, %i8_init1_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_0_1 = vector.mma %i8_lhs1_1, %i8_rhs0_1, %i8_mma1_0_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_0_2 = vector.mma %i8_lhs1_2, %i8_rhs0_2, %i8_mma1_0_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_0_3 = vector.mma %i8_lhs1_3, %i8_rhs0_3, %i8_mma1_0_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init1_1 = vector.fragment %i8_ic1_1 shape [%c16, %c16] : vector<8xi32> + %i8_mma1_1_0 = vector.mma %i8_lhs1_0, %i8_rhs1_0, %i8_init1_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_1_1 = vector.mma %i8_lhs1_1, %i8_rhs1_1, %i8_mma1_1_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_1_2 = vector.mma %i8_lhs1_2, %i8_rhs1_2, %i8_mma1_1_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_1_3 = vector.mma %i8_lhs1_3, %i8_rhs1_3, %i8_mma1_1_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init1_2 = vector.fragment %i8_ic1_2 shape [%c16, %c16] : vector<8xi32> + %i8_mma1_2_0 = vector.mma %i8_lhs1_0, %i8_rhs2_0, %i8_init1_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_2_1 = vector.mma %i8_lhs1_1, %i8_rhs2_1, %i8_mma1_2_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_2_2 = vector.mma %i8_lhs1_2, %i8_rhs2_2, %i8_mma1_2_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_2_3 = vector.mma %i8_lhs1_3, %i8_rhs2_3, %i8_mma1_2_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + %i8_init1_3 = vector.fragment %i8_ic1_3 shape [%c16, %c16] : vector<8xi32> + %i8_mma1_3_0 = vector.mma %i8_lhs1_0, %i8_rhs3_0, %i8_init1_3 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_3_1 = vector.mma %i8_lhs1_1, %i8_rhs3_1, %i8_mma1_3_0 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_3_2 = vector.mma %i8_lhs1_2, %i8_rhs3_2, %i8_mma1_3_1 : vector<4xi32>, vector<4xi32>, vector<8xi32> + %i8_mma1_3_3 = vector.mma %i8_lhs1_3, %i8_rhs3_3, %i8_mma1_3_2 : vector<4xi32>, vector<4xi32>, vector<8xi32> + scf.schedule.fence + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.yield %i8_mma0_0_3, %i8_mma0_1_3, %i8_mma0_2_3, %i8_mma0_3_3, %i8_mma1_0_3, %i8_mma1_1_3, %i8_mma1_2_3, %i8_mma1_3_3 : vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32>, vector<8xi32> + } + %i8_asi0_0 = index.add %asb0, %c0 : index + %i8_asi0_1 = index.mul %i8_asi0_0, %c2 : index + %i8_asi0_2 = index.add %i8_asi0_1, %lane_hi : index + %i8_asi0_3 = index.mul %i8_asi0_2, %c8 : index + %i8_asv0, %i8_unused_sum0 = func.call @ggml_symmetric_i8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %i8_asi0_3) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + %i8_asi1_0 = index.add %asb1, %c0 : index + %i8_asi1_1 = index.mul %i8_asi1_0, %c2 : index + %i8_asi1_2 = index.add %i8_asi1_1, %lane_hi : index + %i8_asi1_3 = index.mul %i8_asi1_2, %c8 : index + %i8_asv1, %i8_unused_sum1 = func.call @ggml_symmetric_i8_wmmai8_load_metadata(%token256, %scratch, %as_off, %asum_off, %i8_asi1_3) : (i1, buffer, offset, offset, index) -> (vector<8xf32>, vector<8xf32>) + %i8_wsi0_0 = index.add %ln0, %lane_lo : index + %i8_wsi0 = index.assume %i8_wsi0_0 [range(%i8_wsi0_0, 0, 63)] : index + %i8_scale_source_row0 = index.add %i8_source_row_base, %i8_wsi0 : index + %i8_scale_byte0_0 = index.mul %i8_scale_source_row0, %c2 : index + %i8_scale_byte0 = index.add %i8_record_base, %i8_scale_byte0_0 : index + %i8_scale_bytes0 = vector.load %i8_weight_view[%i8_scale_byte0] : view<[%i8_weight_bytes]xi8> -> vector<2xi8> + %i8_scale_halfv0 = vector.bitcast %i8_scale_bytes0 : vector<2xi8> to vector<1xf16> + %i8_scale_half0 = vector.extract %i8_scale_halfv0[0] : vector<1xf16> -> f16 + %i8_wscale0_scalar = scalar.extf %i8_scale_half0 : f16 to f32 + %i8_wscale0 = vector.splat %i8_wscale0_scalar : vector<8xf32> + %i8_wsi1_0 = index.add %ln1, %lane_lo : index + %i8_wsi1 = index.assume %i8_wsi1_0 [range(%i8_wsi1_0, 0, 63)] : index + %i8_scale_source_row1 = index.add %i8_source_row_base, %i8_wsi1 : index + %i8_scale_byte1_0 = index.mul %i8_scale_source_row1, %c2 : index + %i8_scale_byte1 = index.add %i8_record_base, %i8_scale_byte1_0 : index + %i8_scale_bytes1 = vector.load %i8_weight_view[%i8_scale_byte1] : view<[%i8_weight_bytes]xi8> -> vector<2xi8> + %i8_scale_halfv1 = vector.bitcast %i8_scale_bytes1 : vector<2xi8> to vector<1xf16> + %i8_scale_half1 = vector.extract %i8_scale_halfv1[0] : vector<1xf16> -> f16 + %i8_wscale1_scalar = scalar.extf %i8_scale_half1 : f16 to f32 + %i8_wscale1 = vector.splat %i8_wscale1_scalar : vector<8xf32> + %i8_wsi2_0 = index.add %ln2, %lane_lo : index + %i8_wsi2 = index.assume %i8_wsi2_0 [range(%i8_wsi2_0, 0, 63)] : index + %i8_scale_source_row2 = index.add %i8_source_row_base, %i8_wsi2 : index + %i8_scale_byte2_0 = index.mul %i8_scale_source_row2, %c2 : index + %i8_scale_byte2 = index.add %i8_record_base, %i8_scale_byte2_0 : index + %i8_scale_bytes2 = vector.load %i8_weight_view[%i8_scale_byte2] : view<[%i8_weight_bytes]xi8> -> vector<2xi8> + %i8_scale_halfv2 = vector.bitcast %i8_scale_bytes2 : vector<2xi8> to vector<1xf16> + %i8_scale_half2 = vector.extract %i8_scale_halfv2[0] : vector<1xf16> -> f16 + %i8_wscale2_scalar = scalar.extf %i8_scale_half2 : f16 to f32 + %i8_wscale2 = vector.splat %i8_wscale2_scalar : vector<8xf32> + %i8_wsi3_0 = index.add %ln3, %lane_lo : index + %i8_wsi3 = index.assume %i8_wsi3_0 [range(%i8_wsi3_0, 0, 63)] : index + %i8_scale_source_row3 = index.add %i8_source_row_base, %i8_wsi3 : index + %i8_scale_byte3_0 = index.mul %i8_scale_source_row3, %c2 : index + %i8_scale_byte3 = index.add %i8_record_base, %i8_scale_byte3_0 : index + %i8_scale_bytes3 = vector.load %i8_weight_view[%i8_scale_byte3] : view<[%i8_weight_bytes]xi8> -> vector<2xi8> + %i8_scale_halfv3 = vector.bitcast %i8_scale_bytes3 : vector<2xi8> to vector<1xf16> + %i8_scale_half3 = vector.extract %i8_scale_halfv3[0] : vector<1xf16> -> f16 + %i8_wscale3_scalar = scalar.extf %i8_scale_half3 : f16 to f32 + %i8_wscale3 = vector.splat %i8_wscale3_scalar : vector<8xf32> + %i8_fp0_0 = vector.sitofp %i8_i0_0 : vector<8xi32> to vector<8xf32> + %i8_scale0_0 = vector.mulf %i8_asv0, %i8_wscale0 : vector<8xf32> + %i8_next0_0 = vector.fmaf %i8_fp0_0, %i8_scale0_0, %fc0_0 : vector<8xf32> + scf.schedule.fence + %i8_fp0_1 = vector.sitofp %i8_i0_1 : vector<8xi32> to vector<8xf32> + %i8_scale0_1 = vector.mulf %i8_asv0, %i8_wscale1 : vector<8xf32> + %i8_next0_1 = vector.fmaf %i8_fp0_1, %i8_scale0_1, %fc0_1 : vector<8xf32> + scf.schedule.fence + %i8_fp0_2 = vector.sitofp %i8_i0_2 : vector<8xi32> to vector<8xf32> + %i8_scale0_2 = vector.mulf %i8_asv0, %i8_wscale2 : vector<8xf32> + %i8_next0_2 = vector.fmaf %i8_fp0_2, %i8_scale0_2, %fc0_2 : vector<8xf32> + scf.schedule.fence + %i8_fp0_3 = vector.sitofp %i8_i0_3 : vector<8xi32> to vector<8xf32> + %i8_scale0_3 = vector.mulf %i8_asv0, %i8_wscale3 : vector<8xf32> + %i8_next0_3 = vector.fmaf %i8_fp0_3, %i8_scale0_3, %fc0_3 : vector<8xf32> + scf.schedule.fence + %i8_fp1_0 = vector.sitofp %i8_i1_0 : vector<8xi32> to vector<8xf32> + %i8_scale1_0 = vector.mulf %i8_asv1, %i8_wscale0 : vector<8xf32> + %i8_next1_0 = vector.fmaf %i8_fp1_0, %i8_scale1_0, %fc1_0 : vector<8xf32> + scf.schedule.fence + %i8_fp1_1 = vector.sitofp %i8_i1_1 : vector<8xi32> to vector<8xf32> + %i8_scale1_1 = vector.mulf %i8_asv1, %i8_wscale1 : vector<8xf32> + %i8_next1_1 = vector.fmaf %i8_fp1_1, %i8_scale1_1, %fc1_1 : vector<8xf32> + scf.schedule.fence + %i8_fp1_2 = vector.sitofp %i8_i1_2 : vector<8xi32> to vector<8xf32> + %i8_scale1_2 = vector.mulf %i8_asv1, %i8_wscale2 : vector<8xf32> + %i8_next1_2 = vector.fmaf %i8_fp1_2, %i8_scale1_2, %fc1_2 : vector<8xf32> + scf.schedule.fence + %i8_fp1_3 = vector.sitofp %i8_i1_3 : vector<8xi32> to vector<8xf32> + %i8_scale1_3 = vector.mulf %i8_asv1, %i8_wscale3 : vector<8xf32> + %i8_next1_3 = vector.fmaf %i8_fp1_3, %i8_scale1_3, %fc1_3 : vector<8xf32> + scf.schedule.fence + scf.yield %i8_next0_0, %i8_next0_1, %i8_next0_2, %i8_next0_3, %i8_next1_0, %i8_next1_1, %i8_next1_2, %i8_next1_3 : vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32>, vector<8xf32> + } + %out0_0 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_0 = vector.fragment.load %gate_view[%gm0, %gn0] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_0 = vector.siluf %gate0_0 : vector<8xf32> + %result0_0 = vector.mulf %silu0_0, %f0_0 : vector<8xf32> + scf.yield %result0_0 : vector<8xf32> + } else { + scf.yield %f0_0 : vector<8xf32> + } + %out0_1 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_1 = vector.fragment.load %gate_view[%gm0, %gn1] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_1 = vector.siluf %gate0_1 : vector<8xf32> + %result0_1 = vector.mulf %silu0_1, %f0_1 : vector<8xf32> + scf.yield %result0_1 : vector<8xf32> + } else { + scf.yield %f0_1 : vector<8xf32> + } + %out0_2 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_2 = vector.fragment.load %gate_view[%gm0, %gn2] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_2 = vector.siluf %gate0_2 : vector<8xf32> + %result0_2 = vector.mulf %silu0_2, %f0_2 : vector<8xf32> + scf.yield %result0_2 : vector<8xf32> + } else { + scf.yield %f0_2 : vector<8xf32> + } + %out0_3 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate0_3 = vector.fragment.load %gate_view[%gm0, %gn3] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu0_3 = vector.siluf %gate0_3 : vector<8xf32> + %result0_3 = vector.mulf %silu0_3, %f0_3 : vector<8xf32> + scf.yield %result0_3 : vector<8xf32> + } else { + scf.yield %f0_3 : vector<8xf32> + } + %out1_0 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_0 = vector.fragment.load %gate_view[%gm1, %gn0] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_0 = vector.siluf %gate1_0 : vector<8xf32> + %result1_0 = vector.mulf %silu1_0, %f1_0 : vector<8xf32> + scf.yield %result1_0 : vector<8xf32> + } else { + scf.yield %f1_0 : vector<8xf32> + } + %out1_1 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_1 = vector.fragment.load %gate_view[%gm1, %gn1] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_1 = vector.siluf %gate1_1 : vector<8xf32> + %result1_1 = vector.mulf %silu1_1, %f1_1 : vector<8xf32> + scf.yield %result1_1 : vector<8xf32> + } else { + scf.yield %f1_1 : vector<8xf32> + } + %out1_2 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_2 = vector.fragment.load %gate_view[%gm1, %gn2] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_2 = vector.siluf %gate1_2 : vector<8xf32> + %result1_2 = vector.mulf %silu1_2, %f1_2 : vector<8xf32> + scf.yield %result1_2 : vector<8xf32> + } else { + scf.yield %f1_2 : vector<8xf32> + } + %out1_3 = scf.if %apply_swiglu -> (vector<8xf32>) { + %gate1_3 = vector.fragment.load %gate_view[%gm1, %gn3] shape [%c16, %c16] : view<[%cols_b]x[%rows_b]xf32> -> vector<8xf32> + %silu1_3 = vector.siluf %gate1_3 : vector<8xf32> + %result1_3 = vector.mulf %silu1_3, %f1_3 : vector<8xf32> + scf.yield %result1_3 : vector<8xf32> + } else { + scf.yield %f1_3 : vector<8xf32> + } + + vector.fragment.store %out0_0, %dst_view[%gm0, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out0_1, %dst_view[%gm0, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out0_2, %dst_view[%gm0, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out0_3, %dst_view[%gm0, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_0, %dst_view[%gm1, %gn0] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_1, %dst_view[%gm1, %gn1] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_2, %dst_view[%gm1, %gn2] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + vector.fragment.store %out1_3, %dst_view[%gm1, %gn3] shape [%c16, %c16] : vector<8xf32>, view<[%cols_b]x[%rows_b]xf32> + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom new file mode 100644 index 000000000000..d94f50c53032 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom @@ -0,0 +1,19 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.add_rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.add_rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon : f32 + +// Publishes residual ADD, RMSNorm, learned scale, and a signed-I4/K64 alternate. +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_add_rmsnorm_binary_symmetric_i4_k32") @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 16)] { + template.apply<@ggml.add_rmsnorm_binary_symmetric_i4_k32.body>(%token_count, %lhs, %rhs, %residual_output, %weight, %output, %i4_qs, %i4_ds, %i4_sums) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32.loom new file mode 100644 index 000000000000..8977d392a703 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32.loom @@ -0,0 +1,42 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_f32.body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32") @ggml_rmsnorm_binary_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { + %publish_f16 = scalar.constant false : i1 + template.apply<@ggml.rmsnorm_binary_f32.body>(%publish_f16, %token_count, %input, %rhs, %output, %output) : (i1, index, buffer, buffer, buffer, buffer) + kernel.return +} + +check.case public @ggml_rmsnorm_binary_f32_mul_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %rhs = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> + check.return +} + +check.case public @ggml_rmsnorm_binary_f32_add_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %rhs = check.generate.iota offset(0.0) step(0.015625) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.iota offset(1.0) step(0.015625) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_f16.loom new file mode 100644 index 000000000000..5344c6590906 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_f16.loom @@ -0,0 +1,20 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_f32.body(%publish_f16: i1, %token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32_f16") @ggml_rmsnorm_binary_f32_f16(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) where [range(%token_count, 1, 1048576)] { + %publish_f16 = scalar.constant true : i1 + template.apply<@ggml.rmsnorm_binary_f32.body>(%publish_f16, %token_count, %input, %rhs, %output, %f16_output) : (i1, index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_k16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_k16.loom new file mode 100644 index 000000000000..5b31244d23e8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_k16.loom @@ -0,0 +1,20 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_f32_k16.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +// Retain two original lane streams per thread and publish both F32 and K16 F16. +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_f32_k16") @ggml_rmsnorm_binary_f32_k16(%token_count: index) { + %c1 = index.constant 1 : index + %c128 = index.constant 128 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c128, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer, %f16_output: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_binary_f32_k16.body>(%token_count, %input, %rhs, %output, %f16_output) : (index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom new file mode 100644 index 000000000000..655cd00ea03b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_f32_lowtoken_striped.loom @@ -0,0 +1,27 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_f32.lowtoken_striped_body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_lowtoken_striped_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_rmsnorm_binary_lowtoken_striped_gfx11_wave32) export("ggml_rmsnorm_binary_f32_lowtoken_striped") @ggml_rmsnorm_binary_f32_lowtoken_striped(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + %c1023 = index.constant 1023 : index + %c1024 = index.constant 1024 : index + %hidden_size = config.get @ggml.rmsnorm_binary_f32.hidden_size : index + %rounded_hidden = index.add %hidden_size, %c1023 : index + %stripe_count = index.div %rounded_hidden, %c1024 : index + %group_count = index.mul %token_count, %stripe_count : index + kernel.launch.config workgroups(%group_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 5)] { + template.apply<@ggml.rmsnorm_binary_f32.lowtoken_striped_body>(%token_count, %input, %rhs, %output) : (index, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom new file mode 100644 index 000000000000..a68764d6b80e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_publish_q8_1_x4_configured.loom @@ -0,0 +1,56 @@ +// 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)] + +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_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 + %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_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_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> + %rhs = check.generate.iota offset(-1.0) step(0.015625) : tensor<128xf32> + %f32 = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.fill value(0) : tensor<144xi8> + %actual = check.generate.fill value(1) : tensor<144xi8> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %f32) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %f32, %expected) : [index, index](index, index, tensor<128xf32>, tensor<144xi8>) + %actual_f32 = check.generate.fill value(0.0) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_binary_q8_1_x4_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_strided_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_strided_f32.loom new file mode 100644 index 000000000000..3ae10cd93821 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_strided_f32.loom @@ -0,0 +1,34 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_strided_f32.body(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_binary_f32.input_stride : %value: index where [range(%value, 128, 1048576)] + +config.decl @ggml.rmsnorm_binary_f32.rms_epsilon : f32 + +config.decl @ggml.rmsnorm_binary_f32.op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_strided_f32") @ggml_rmsnorm_binary_strided_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_binary_strided_f32.body>(%token_count, %input, %rhs, %output) : (index, buffer, buffer, buffer) + kernel.return +} + +check.case public @ggml_rmsnorm_binary_strided_f32_mul_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(2.0) : tensor<2x160xf32> + %rhs = check.generate.fill value(3.0) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<2x128xf32> + %expected = check.generate.fill value(3.0) : tensor<2x128xf32> + kernel.launch @ggml_rmsnorm_binary_strided_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<2x160xf32>, tensor<128xf32>, tensor<2x128xf32>) + check.expect.close actual(%output) expected(%expected) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x128xf32> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom new file mode 100644 index 000000000000..4735fe61fc98 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom @@ -0,0 +1,18 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_binary_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_binary_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_binary_symmetric_i4.rms_epsilon : f32 + +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_binary_symmetric_i4_k32") @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 16)] { + template.apply<@ggml.rmsnorm_binary_symmetric_i4_k32.body>(%token_count, %input, %weight, %output, %i4_qs, %i4_ds, %i4_sums) : (index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_f32.loom new file mode 100644 index 000000000000..fed98e520868 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_f32.loom @@ -0,0 +1,29 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.body(%token_count: index, %input: buffer, %output: buffer) + +amdgpu.target @ggml_rmsnorm_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_f32.hidden_size : %value: index where [range(%value, 64, 32768), mul(%value, 64)] +config.decl @ggml.rmsnorm_f32.input_stride : %value: index where [range(%value, 64, 1048576)] +config.decl @ggml.rmsnorm_f32.rms_epsilon : f32 + +kernel.def target(@ggml_rmsnorm_gfx11_wave32) export("ggml_rmsnorm_f32") @ggml_rmsnorm_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rmsnorm_f32.body>(%token_count, %input, %output) : (index, buffer, buffer) + kernel.return +} + +check.case public @ggml_rmsnorm_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.fill value(1.0) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_f32[%token_count](%token_count, %input, %output) : [index](index, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999995e-07) rtol(9.9999999999999995e-07) nan(same) : tensor<128xf32> + check.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish.loom new file mode 100644 index 000000000000..64da5ef0bf73 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_publish.loom @@ -0,0 +1,32 @@ +// 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)] + +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 + %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, %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_silu_mul_symmetric_i4_k32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom new file mode 100644 index 000000000000..331f66a9849d --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom @@ -0,0 +1,19 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) + +amdgpu.target @ggml_rmsnorm_binary_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon : f32 + +// Publishes RMSNorm, learned scale, SiLU gate product, and a signed-I4/K64 alternate. +kernel.def target(@ggml_rmsnorm_binary_gfx11_wave32) export("ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32") @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %i4_qs: buffer, %i4_ds: buffer, %i4_sums: buffer) where [range(%token_count, 1, 1048576)] { + template.apply<@ggml.rmsnorm_gate_silu_mul_symmetric_i4_k32.body>(%token_count, %input, %weight, %raw_gate, %output, %i4_qs, %i4_ds, %i4_sums) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_add_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_add_f32.loom new file mode 100644 index 000000000000..55ccbe655c28 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_add_f32.loom @@ -0,0 +1,69 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_f32.row_scale(%token_count: index, %token: index, %hidden_size: index, %input_stride: index, %epsilon: f32, %input: buffer) -> (f32, index, index, i1) + +template.decl @ggml.rmsnorm_f32.apply(%value: f32, %scale: f32) -> (f32) +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_rmsnorm_mul_add_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rmsnorm_mul_add_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] + +config.decl @ggml.rmsnorm_mul_add_f32.rms_epsilon : f32 +config.decl @ggml.rmsnorm_mul_add_f32.normalized_op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.rmsnorm_mul_add_f32.output_op : %value: index where [range(%value, 0, 3)] +config.decl @ggml.rmsnorm_mul_add_f32.binary_lhs : %value: index where [range(%value, 0, 1)] + +kernel.def target(@ggml_rmsnorm_mul_add_gfx11_wave32) export("ggml_rmsnorm_mul_add_f32") @ggml_rmsnorm_mul_add_f32(%token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %input: buffer, %weight: buffer, %residual: buffer, %output: buffer) where [range(%token_count, 1, 1048576)] { + %hidden0 = config.get @ggml.rmsnorm_mul_add_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 32768), mul(%hidden0, 128)] : index + %epsilon = config.get @ggml.rmsnorm_mul_add_f32.rms_epsilon : f32 + %normalized_op = config.get @ggml.rmsnorm_mul_add_f32.normalized_op : index + %output_op = config.get @ggml.rmsnorm_mul_add_f32.output_op : index + %binary_lhs = config.get @ggml.rmsnorm_mul_add_f32.binary_lhs : index + %token0 = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %c0_offset = index.constant 0 : offset + %c256 = index.constant 256 : index + %input_noalias, %weight_noalias, %residual_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %residual, %output : buffer, buffer, buffer, buffer + %row_scale, %token, %launch_token_count, %valid_token = template.apply<@ggml.rmsnorm_f32.row_scale>(%token_count, %token0, %hidden, %hidden, %epsilon, %input_noalias) : (index, index, index, index, f32, buffer) -> (f32, index, index, i1) + %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden]xf32> + %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden]xf32> + %residual_view = buffer.view %residual_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden]xf32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden]xf32> + scf.if %valid_token { + scf.for %channel = [%workitem to %hidden step %c256] { + %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden]xf32> -> f32 + %weight_value = view.load %weight_view[%channel] : view<[%hidden]xf32> -> f32 + %residual_value = view.load %residual_view[%token, %channel] : view<[%launch_token_count]x[%hidden]xf32> -> f32 + %normalized = template.apply<@ggml.rmsnorm_f32.apply>(%value, %row_scale) : (f32, f32) -> (f32) + %scaled = template.apply<@ggml.binary_f32.apply>(%normalized_op, %normalized, %weight_value) : (index, f32, f32) -> (f32) + %c1 = index.constant 1 : index + %is_binary_lhs = index.cmp eq, %binary_lhs, %c1 : index + %lhs = scf.select %is_binary_lhs, %scaled, %residual_value : f32 + %rhs = scf.select %is_binary_lhs, %residual_value, %scaled : f32 + %result = template.apply<@ggml.binary_f32.apply>(%output_op, %lhs, %rhs) : (index, f32, f32) -> (f32) + view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%hidden]xf32> + } + } + kernel.return +} + +check.case public @ggml_rmsnorm_mul_add_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<128xf32> + %weight = check.generate.iota offset(0.25) step(0.01) : tensor<128xf32> + %residual = check.generate.fill value(3.0) : tensor<128xf32> + %output = check.generate.fill value(0.0) : tensor<128xf32> + %expected = check.generate.iota offset(3.25) step(0.01) : tensor<128xf32> + kernel.launch @ggml_rmsnorm_mul_add_f32[%token_count](%token_count, %input, %weight, %residual, %output) : [index](index, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>, tensor<128xf32>) + check.expect.close actual(%output) expected(%expected) atol(9.9999999999999991e-06) rtol(9.9999999999999991e-06) nan(same) : tensor<128xf32> + check.return +} + +check.benchmark<@ggml_rmsnorm_mul_add_f32_case> @ggml_rmsnorm_mul_add_f32_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom new file mode 100644 index 000000000000..e12c1dc82681 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/normalization/rmsnorm_mul_rope_f32.loom @@ -0,0 +1,43 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.rmsnorm_mul_rope_f32.body(%input: buffer, %weight: buffer, %positions: buffer, %output: buffer) + +amdgpu.target @ggml_rmsnorm_gfx11_wave32 {subgroup_size = 32} + +// One workgroup preserves each strided RMSNorm row before scale and RoPE. +config.decl @ggml.rmsnorm_mul_rope.hidden_size : %value: index where [range(%value, 2, 512)] +config.decl @ggml.rmsnorm_mul_rope.ne1 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.rmsnorm_mul_rope.ne2 : %value: index where [range(%value, 1, 1048576)] +config.decl @ggml.rmsnorm_mul_rope.ne3 : %value: index where [range(%value, 1, 1024)] +config.decl @ggml.rmsnorm_mul_rope.input_stride1 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.input_stride2 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.input_stride3 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.output_stride1 : %value: index where [range(%value, 0, 268435456)] +config.decl @ggml.rmsnorm_mul_rope.output_stride2 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.output_stride3 : %value: index where [range(%value, 0, 1073741823)] +config.decl @ggml.rmsnorm_mul_rope.n_dims : %value: index where [range(%value, 2, 512)] +config.decl @ggml.rmsnorm_mul_rope.section0 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.section1 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.section2 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.section3 : %value: index where [range(%value, 0, 65536)] +config.decl @ggml.rmsnorm_mul_rope.mode : %value: index where [range(%value, 2, 40)] +config.decl @ggml.rmsnorm_mul_rope.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] +config.decl @ggml.rmsnorm_mul_rope.epsilon : f32 +config.decl @ggml.rmsnorm_mul_rope.freq_base : f32 +config.decl @ggml.rmsnorm_mul_rope.freq_scale : f32 +config.decl @ggml.rmsnorm_mul_rope.attn_factor : f32 + +kernel.def target(@ggml_rmsnorm_gfx11_wave32) export("ggml_rmsnorm_mul_rope_f32") @ggml_rmsnorm_mul_rope_f32() { + %unit = index.constant 1 : index + %ne1 = config.get @ggml.rmsnorm_mul_rope.ne1 : index + %ne2 = config.get @ggml.rmsnorm_mul_rope.ne2 : index + %ne3 = config.get @ggml.rmsnorm_mul_rope.ne3 : index + %wg = config.get @ggml.rmsnorm_mul_rope.workgroup_size : index + %n12 = index.mul %ne1, %ne2 : index + %rows = index.mul %n12, %ne3 : index + kernel.launch.config workgroups(%rows, %unit, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%input: buffer, %weight: buffer, %positions: buffer, %output: buffer) { + template.apply<@ggml.rmsnorm_mul_rope_f32.body>(%input, %weight, %positions, %output) : (buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_concat_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_concat_f32.loom new file mode 100644 index 000000000000..b7e03b4ab1ad --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_concat_f32.loom @@ -0,0 +1,114 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Publishes an unrotated prefix and a partially rotated suffix directly into +// one packed head. The source remains in its original strided full-head layout. +template.decl @ggml.rope_f32.rotate_neox_packet(%position: f32, %input_base: index, %pair_channel: index, %half_n_dims: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) +template.decl @ggml.rope_f32.rotate_normal_packet(%position: f32, %input_base: index, %normal_channel: index, %bounded_input_span: index, %mscale: f32, %input: buffer, %theta_packet: vector<2xf32>, %freq_factors_packet: vector<2xf32>) -> (vector<4xf32>) +template.decl @ggml.rope_f32.copy_packet(%input_base: index, %normal_channel: index, %bounded_input_span: index, %input: buffer) -> (vector<4xf32>) + +amdgpu.target @ggml_rope_concat_f32_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rope_concat_f32.prefix_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_concat_f32.rope_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_concat_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_concat_f32.head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.rope_concat_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.rope_concat_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_concat_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_concat_f32.input_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.rope_concat_f32.mscale : f32 +config.decl @ggml.rope_concat_f32.mode : %value: index where [range(%value, 0, 2)] + +kernel.def target(@ggml_rope_concat_f32_gfx11_wave32) export("ggml_rope_concat_f32") @ggml_rope_concat_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_concat_f32.token_capacity : index + %head_count = config.get @ggml.rope_concat_f32.head_count : index + %prefix_size = config.get @ggml.rope_concat_f32.prefix_size : index + %rope_size = config.get @ggml.rope_concat_f32.rope_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %output_size = index.add %prefix_size, %rope_size : index + %packet_count = index.div %output_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%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)] { + %token_capacity = config.get @ggml.rope_concat_f32.token_capacity : index + %head_count0 = config.get @ggml.rope_concat_f32.head_count : index + %prefix_size0 = config.get @ggml.rope_concat_f32.prefix_size : index + %rope_size0 = config.get @ggml.rope_concat_f32.rope_size : index + %n_dims0 = config.get @ggml.rope_concat_f32.n_dims : index + %input_stride1_0 = config.get @ggml.rope_concat_f32.input_stride1 : index + %input_stride2_0 = config.get @ggml.rope_concat_f32.input_stride2 : index + %input_span0 = config.get @ggml.rope_concat_f32.input_span : index + %mscale = config.get @ggml.rope_concat_f32.mscale : f32 + %mode = config.get @ggml.rope_concat_f32.mode : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_input_span = index.assume %input_span0 [range(%input_span0, 1, 1073741824)] : index + %head_count, %prefix_size, %rope_size, %n_dims = index.assume %head_count0, %prefix_size0, %rope_size0, %n_dims0 [range(%head_count0, 1, 64), range(%prefix_size0, 4, 1024), range(%rope_size0, 4, 1024), range(%n_dims0, 4, 1024), mul(%prefix_size0, 4), mul(%rope_size0, 4), mul(%n_dims0, 4), le(%n_dims0, %rope_size0)] : index, 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 + %output_size0 = index.add %prefix_size, %rope_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 8, 1024), mul(%output_size0, 4)] : index + %head = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %packet0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %mode_neox = index.constant 2 : index + %c0_offset = index.constant 0 : offset + %prefix_packet_count = index.div %prefix_size, %c4 : index + %rotary_packet_count = index.div %n_dims, %c4 : index + %packet_count = index.div %output_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %active_packet = index.cmp ult, %packet0, %packet_count : index + %publish = scalar.andi %valid_token, %active_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %is_prefix = index.cmp ult, %packet0, %prefix_packet_count : index + %rope_packet0 = index.sub %packet0, %prefix_packet_count : index + %rope_packet = scf.select %is_prefix, %c0, %rope_packet0 : index + %rotary_packet = index.cmp ult, %rope_packet, %rotary_packet_count : index + %tail_packet = index.cmp uge, %rope_packet, %rotary_packet_count : index + %is_neox = index.cmp eq, %mode, %mode_neox : index + %positions_noalias, %input_noalias, %theta_noalias, %freq_factors_noalias, %output_noalias = buffer.assume.noalias %positions, %input, %theta, %freq_factors, %output : buffer, buffer, buffer, buffer, buffer + scf.if %publish { + %packet = index.assume %packet0 [lt(%packet0, %packet_count)] : index + %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + %input_token_offset = index.mul %token, %input_stride2 : index + %input_head_offset = index.mul %head, %input_stride1 : index + %input_base = index.add %input_token_offset, %input_head_offset : index + %output_channel = index.mul %packet, %c4 : index + scf.if %is_prefix { + %values = template.apply<@ggml.rope_f32.copy_packet>(%input_base, %output_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + vector.store %values, %output_view[%token, %head, %output_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } else { + %rope_input_base = index.add %input_base, %prefix_size : index + %rope_channel = index.mul %rope_packet, %c4 : index + scf.if %rotary_packet { + %half_n_dims = index.div %n_dims, %c2 : index + %pair_channel = index.mul %rope_packet, %c2 : index + %theta_view = buffer.view %theta_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %freq_factors_view = buffer.view %freq_factors_noalias[%c0_offset] : buffer -> view<[%half_n_dims]xf32> + %position_i32 = view.load %positions_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %position = scalar.sitofp %position_i32 : i32 to f32 + %theta_packet = vector.load %theta_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + %freq_factors_packet = vector.load %freq_factors_view[%pair_channel] : view<[%half_n_dims]xf32> -> vector<2xf32> + scf.if %is_neox { + %paired_channel = index.add %pair_channel, %half_n_dims : index + %scaled_low, %scaled_high = template.apply<@ggml.rope_f32.rotate_neox_packet>(%position, %rope_input_base, %pair_channel, %half_n_dims, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) + %low_output_channel = index.add %prefix_size, %pair_channel : index + %high_output_channel = index.add %prefix_size, %paired_channel : index + vector.store %scaled_low, %output_view[%token, %head, %low_output_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + vector.store %scaled_high, %output_view[%token, %head, %high_output_channel] : vector<2xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } else { + %rotated = template.apply<@ggml.rope_f32.rotate_normal_packet>(%position, %rope_input_base, %rope_channel, %bounded_input_span, %mscale, %input_noalias, %theta_packet, %freq_factors_packet) : (f32, index, index, index, f32, buffer, vector<2xf32>, vector<2xf32>) -> (vector<4xf32>) + vector.store %rotated, %output_view[%token, %head, %output_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } + } + scf.if %tail_packet { + %values = template.apply<@ggml.rope_f32.copy_packet>(%rope_input_base, %rope_channel, %bounded_input_span, %input_noalias) : (index, index, index, buffer) -> (vector<4xf32>) + vector.store %values, %output_view[%token, %head, %output_channel] : vector<4xf32>, view<[%bounded_token_count]x[%head_count]x[%output_size]xf32> + } + } + } + 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 new file mode 100644 index 000000000000..b321fff4e636 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_f32.loom @@ -0,0 +1,30 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Generic GGML ROPE for F32 head-major rows. +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} + +config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.rope_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_f32.input_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.rope_f32.mscale : f32 +config.decl @ggml.rope_f32.mode : %value: index where [range(%value, 0, 2)] + +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 + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %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/rope/rope_scale_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_scale_f32.loom new file mode 100644 index 000000000000..0f059de07ff7 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_scale_f32.loom @@ -0,0 +1,37 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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) + +amdgpu.target @ggml_rope_scale_f32_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rope_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_f32.head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.rope_f32.token_capacity : %value: index where [range(%value, 1, 2048)] + +kernel.def target(@ggml_rope_scale_f32_gfx11_wave32) export("ggml_rope_scale_f32") @ggml_rope_scale_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %head_count = config.get @ggml.rope_f32.head_count : index + %head_size = config.get @ggml.rope_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_scale_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %output) : (index, buffer, buffer, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@ggml_rope_scale_f32_gfx11_wave32) export("ggml_rope_token_scale_f32") @ggml_rope_token_scale_f32(%token_count: index) { + %token_capacity = config.get @ggml.rope_f32.token_capacity : index + %head_count = config.get @ggml.rope_f32.head_count : index + %head_size = config.get @ggml.rope_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %positions: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %scales: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_token_scale_f32.body>(%token_count, %positions, %input, %theta, %freq_factors, %scales, %output) : (index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_set_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_set_rows_f32.loom new file mode 100644 index 000000000000..61f7bdc2baaa --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rope/rope_set_rows_f32.loom @@ -0,0 +1,30 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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) + +amdgpu.target @ggml_rope_set_rows_f32_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.rope_set_rows_f32.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_set_rows_f32.n_dims : %value: index where [range(%value, 4, 1024), mul(%value, 4)] +config.decl @ggml.rope_set_rows_f32.head_count : %value: index where [range(%value, 1, 64)] +config.decl @ggml.rope_set_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.rope_set_rows_f32.input_stride1 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_set_rows_f32.input_stride2 : %value: index where [range(%value, 4, 1073741824)] +config.decl @ggml.rope_set_rows_f32.input_span : %value: index where [range(%value, 1, 1073741824)] +config.decl @ggml.rope_set_rows_f32.mscale : f32 +config.decl @ggml.rope_set_rows_f32.output_format : %value: index where [range(%value, 16, 32)] +config.decl @ggml.rope_set_rows_f32.mode : %value: index where [range(%value, 0, 2)] + +kernel.def target(@ggml_rope_set_rows_f32_gfx11_wave32) @ggml_rope_set_rows_f32(%token_count: index, %cache_row_count: index) { + %token_capacity = config.get @ggml.rope_set_rows_f32.token_capacity : index + %head_count = config.get @ggml.rope_set_rows_f32.head_count : index + %head_size = config.get @ggml.rope_set_rows_f32.head_size : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %pair_packet_count = index.div %head_size, %c4 : index + kernel.launch.config workgroups(%head_count, %token_capacity, %c1) workgroup_size(%pair_packet_count, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %positions: buffer, %indices: buffer, %input: buffer, %theta: buffer, %freq_factors: buffer, %cache: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.rope_set_rows_f32.body>(%token_count, %cache_row_count, %positions, %indices, %input, %theta, %freq_factors, %cache) : (index, index, buffer, buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/gather_add_rmsnorm_binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/gather_add_rmsnorm_binary_f32.loom new file mode 100644 index 000000000000..7434e3648561 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/gather_add_rmsnorm_binary_f32.loom @@ -0,0 +1,133 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +// Publishes both the indexed residual and its RMS-normalized learned-scale +// output. Each output row may select any valid source row independently. +template.decl @ggml.binary_f32.apply(%op: index, %lhs: f32, %rhs: f32) -> (f32) + +amdgpu.target @ggml_gather_add_rmsnorm_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.gather_add_rmsnorm_binary_f32.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.gather_add_rmsnorm_binary_f32.rms_epsilon : f32 +config.decl @ggml.gather_add_rmsnorm_binary_f32.gather_op : %value: index where [range(%value, 0, 3)] + +kernel.def target(@ggml_gather_add_rmsnorm_gfx11_wave32) export("ggml_gather_add_rmsnorm_binary_f32") @ggml_gather_add_rmsnorm_binary_f32(%source_token_count: index, %output_token_count: index) { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%output_token_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%source_token_count: index, %output_token_count: index, %attention: buffer, %residual: buffer, %output_ids: buffer, %raw_output: buffer, %weight: buffer, %normalized_output: buffer) where [range(%source_token_count, 1, 2048), range(%output_token_count, 1, 2048)] { + %source_count = index.assume %source_token_count [range(%source_token_count, 1, 2048)] : index + %output_count = index.assume %output_token_count [range(%output_token_count, 1, 2048)] : index + %hidden0 = config.get @ggml.gather_add_rmsnorm_binary_f32.hidden_size : index + %hidden = index.assume %hidden0 [range(%hidden0, 128, 32768), mul(%hidden0, 128)] : index + %epsilon = config.get @ggml.gather_add_rmsnorm_binary_f32.rms_epsilon : f32 + %gather_op = config.get @ggml.gather_add_rmsnorm_binary_f32.gather_op : index + %output_row0 = kernel.workgroup.id : index + %output_row = index.assume %output_row0 [lt(%output_row0, %output_count)] : index + %workitem = kernel.workitem.id : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index + %lane = kernel.subgroup.lane.id : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %base = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %scratch_bytes = index.constant 1024 : offset + + %attention_noalias, %residual_noalias, %ids_noalias, %raw_noalias, %weight_noalias, %normalized_noalias = buffer.assume.noalias %attention, %residual, %output_ids, %raw_output, %weight, %normalized_output : buffer, buffer, buffer, buffer, buffer, buffer + %attention_view = buffer.view %attention_noalias[%base] : buffer -> view<[%source_count]x[%hidden]xf32> + %residual_view = buffer.view %residual_noalias[%base] : buffer -> view<[%source_count]x[%hidden]xf32> + %ids_view = buffer.view %ids_noalias[%base] : buffer -> view<[%output_count]xi32> + %raw_view = buffer.view %raw_noalias[%base] : buffer -> view<[%output_count]x[%hidden]xf32> + %weight_view = buffer.view %weight_noalias[%base] : buffer -> view<[%hidden]xf32> + %normalized_view = buffer.view %normalized_noalias[%base] : buffer -> view<[%output_count]x[%hidden]xf32> + %source_row_i32 = view.load %ids_view[%output_row] : view<[%output_count]xi32> -> i32 + %source_row0 = index.cast %source_row_i32 : i32 to index + %source_row = index.assume %source_row0 [range(%source_row0, 0, 2047), lt(%source_row0, %source_count)] : index + + %thread_sum = scf.for %channel = [%workitem to %hidden step %c256](%running_sum = %c0_f32 : f32) -> (f32) { + %attention_value = view.load %attention_view[%source_row, %channel] : view<[%source_count]x[%hidden]xf32> -> f32 + %residual_value = view.load %residual_view[%source_row, %channel] : view<[%source_count]x[%hidden]xf32> -> f32 + %combined = template.apply<@ggml.binary_f32.apply>(%gather_op, %attention_value, %residual_value) : (index, f32, f32) -> (f32) + view.store %combined, %raw_view[%output_row, %channel] : f32, view<[%output_count]x[%hidden]xf32> + %square = scalar.mulf %combined, %combined : f32 + %next_sum = scalar.addf %running_sum, %square : f32 + scf.yield %next_sum : f32 + } + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %scratch = buffer.alloca align(16) %scratch_bytes : buffer + %scratch_view = buffer.view %scratch[%base] : buffer -> view<256xf32> + %is_subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %is_subgroup_leader { + view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index + %is_reduction_lane = index.cmp ult, %lane, %c8 : index + %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 + %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { + %value = view.load %scratch_view[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 + %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 + scf.if %writes_scale { + %hidden_i32 = index.cast %hidden : index to i32 + %hidden_f32 = scalar.sitofp %hidden_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_view[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %row_scale = view.load %scratch_view[%c0] : view<256xf32> -> f32 + scf.for %channel = [%workitem to %hidden step %c256] { + %value = view.load %raw_view[%output_row, %channel] : view<[%output_count]x[%hidden]xf32> -> f32 + %learned_weight = view.load %weight_view[%channel] : view<[%hidden]xf32> -> f32 + %normalized = scalar.mulf %value, %row_scale : f32 + %result = scalar.mulf %normalized, %learned_weight : f32 + view.store %result, %normalized_view[%output_row, %channel] : f32, view<[%output_count]x[%hidden]xf32> + } + kernel.return +} + +check.case public @ggml_gather_add_rmsnorm_binary_f32_index_case { + %source_count = check.literal value(3) : index + %output_count = check.literal value(3) : index + %attention = check.generate.fill value(1.0) : tensor<3x128xf32> + %residual = check.generate.fill value(3.0) : tensor<3x128xf32> + %ids = check.generate.iota offset(2) step(-1) : tensor<3xi32> + %raw = check.generate.fill value(0.0) : tensor<3x128xf32> + %weight = check.generate.fill value(2.0) : tensor<128xf32> + %normalized = check.generate.fill value(0.0) : tensor<3x128xf32> + %expected_raw = check.generate.fill value(4.0) : tensor<3x128xf32> + %expected_normalized = check.generate.fill value(2.0) : tensor<3x128xf32> + kernel.launch @ggml_gather_add_rmsnorm_binary_f32[%source_count, %output_count](%source_count, %output_count, %attention, %residual, %ids, %raw, %weight, %normalized) : [index, index](index, index, tensor<3x128xf32>, tensor<3x128xf32>, tensor<3xi32>, tensor<3x128xf32>, tensor<128xf32>, tensor<3x128xf32>) + check.expect.close actual(%raw) expected(%expected_raw) atol(0.0) rtol(0.0) nan(same) : tensor<3x128xf32> + check.expect.close actual(%normalized) expected(%expected_normalized) atol(0.000001) rtol(0.000001) nan(same) : tensor<3x128xf32> + check.return +} + +check.benchmark<@ggml_gather_add_rmsnorm_binary_f32_index_case> @ggml_gather_add_rmsnorm_binary_f32_index + +check.case public @ggml_gather_add_rmsnorm_binary_f32_repeated_index_case { + %source_count = check.literal value(3) : index + %output_count = check.literal value(3) : index + %attention = check.generate.fill value(1.0) : tensor<3x128xf32> + %residual = check.generate.fill value(3.0) : tensor<3x128xf32> + %ids = check.generate.fill value(2) : tensor<3xi32> + %raw = check.generate.fill value(0.0) : tensor<3x128xf32> + %weight = check.generate.fill value(2.0) : tensor<128xf32> + %normalized = check.generate.fill value(0.0) : tensor<3x128xf32> + %expected_raw = check.generate.fill value(4.0) : tensor<3x128xf32> + %expected_normalized = check.generate.fill value(2.0) : tensor<3x128xf32> + kernel.launch @ggml_gather_add_rmsnorm_binary_f32[%source_count, %output_count](%source_count, %output_count, %attention, %residual, %ids, %raw, %weight, %normalized) : [index, index](index, index, tensor<3x128xf32>, tensor<3x128xf32>, tensor<3xi32>, tensor<3x128xf32>, tensor<128xf32>, tensor<3x128xf32>) + check.expect.close actual(%raw) expected(%expected_raw) atol(0.0) rtol(0.0) nan(same) : tensor<3x128xf32> + check.expect.close actual(%normalized) expected(%expected_normalized) atol(0.000001) rtol(0.000001) nan(same) : tensor<3x128xf32> + check.return +} + +check.benchmark<@ggml_gather_add_rmsnorm_binary_f32_repeated_index_case> @ggml_gather_add_rmsnorm_binary_f32_repeated_index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_f32.loom new file mode 100644 index 000000000000..a748dfbbce72 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_f32.loom @@ -0,0 +1,60 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.get_rows_f32.body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) + +amdgpu.target @ggml_get_rows_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.get_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.get_rows_f32.hidden_capacity : %value: index where [range(%value, 4, 1073741824), mul(%value, 4)] +config.decl @ggml.get_rows_f32.weight_format : %value: index where [range(%value, 4, 81)] + +kernel.def target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %workgroup_size = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %workgroup_size : index + kernel.launch.config workgroups(%packet_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.get_rows_f32.body>(%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, buffer, buffer, buffer) + kernel.return +} + +// Uniform 0x55 Q4_K bytes encode d=dmin=85.3125, scale=minimum=21, +// and q=5 in every group, producing 85.3125 * 21 * (5 - 1) = 7166.25. +check.case public @ggml_get_rows_f32_q4k_decode_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(1) : index + %hidden_size = check.literal value(2048) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(85) : tensor<1x8x144xi8> + %output = check.generate.fill value(0.0) : tensor<1x2048xf32> + %expected = check.generate.fill value(7166.25) : tensor<1x2048xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<1x8x144xi8>, tensor<1x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> + check.return +} + +// Reverse-order IDs exercise the first, interior, and final legal row while +// remaining suitable for access-sanitized execution. +check.case public @ggml_get_rows_f32_q4k_row_access_case { + %token_count = check.literal value(3) : index + %row_count = check.literal value(3) : index + %hidden_size = check.literal value(2048) : index + %token_ids = check.generate.iota offset(2) step(-1) : tensor<3xi32> + %weight = check.generate.fill value(0) : tensor<3x8x144xi8> + %output = check.generate.fill value(1.0) : tensor<3x2048xf32> + %expected = check.generate.fill value(0.0) : tensor<3x2048xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<3xi32>, tensor<3x8x144xi8>, tensor<3x2048xf32>) + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<3x2048xf32> + check.return +} + +check.benchmark<@ggml_get_rows_f32_q4k_decode_case> @ggml_get_rows_f32_q4k_decode + +check.benchmark<@ggml_get_rows_f32_q4k_row_access_case> @ggml_get_rows_f32_q4k_row_access diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_next_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_next_f32.loom new file mode 100644 index 000000000000..9ee42effa214 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_next_f32.loom @@ -0,0 +1,27 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.get_rows_f32.next_body(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer, %next_output: buffer) + +amdgpu.target @ggml_get_rows_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.get_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.get_rows_f32.hidden_capacity : %value: index where [range(%value, 4, 1073741824), mul(%value, 4)] +config.decl @ggml.get_rows_f32.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.get_rows_f32.next_format : %value: index where [range(%value, 16, 81)] + +kernel.def target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32_next(%token_count: index, %row_count: index, %hidden_size: index) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %workgroup_size = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %workgroup_size : index + kernel.launch.config workgroups(%packet_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer, %next_output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.get_rows_f32.next_body>(%token_count, %row_count, %hidden_size, %token_ids, %weight, %output, %next_output) : (index, index, index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom new file mode 100644 index 000000000000..c0fa08684b8f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_rmsnorm_binary_q8_1_x4_f16.loom @@ -0,0 +1,111 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.get_rows_f32.load_vector4_common(%weight_format: index, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) +template.decl @ggml.normalization.publish_q8.binary_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_get_rows_rmsnorm_gfx11_wave32 {subgroup_size = 32} + +config.decl @ggml.get_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.get_rows_f32.hidden_capacity : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +config.decl @ggml.get_rows_f32.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.get_rows_rmsnorm.rms_epsilon : f32 + +kernel.def target(@ggml_get_rows_rmsnorm_gfx11_wave32) export("ggml_get_rows_rmsnorm_binary_q8_1_x4_f16") @ggml_get_rows_rmsnorm_binary_q8_1_x4_f16(%token_count: index, %row_count: index, %hidden_size: index) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%token_capacity, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %norm_weight: buffer, %raw_output: buffer, %output: buffer, %q8_output: buffer, %f16_output: buffer) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %weight_format = config.get @ggml.get_rows_f32.weight_format : index + %epsilon = config.get @ggml.get_rows_rmsnorm.rms_epsilon : f32 + %bounded_token_count, %configured_token_capacity = index.assume %token_count, %token_capacity [range(%token_count, 1, 2048), eq(%token_count, %token_capacity)] : index, index + %bounded_row_count = index.assume %row_count [range(%row_count, 1, 262208)] : index + %bounded_hidden_size, %configured_hidden_capacity = index.assume %hidden_size, %hidden_capacity [range(%hidden_size, 128, 32768), mul(%hidden_size, 128), eq(%hidden_size, %hidden_capacity)] : index, index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %workitem0 = kernel.workitem.id : index + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %subgroup0 = kernel.subgroup.id : index + %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : 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 + %c128 = index.constant 128 : index + %c1024 = index.constant 1024 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_f32 = scalar.constant 0.0 : f32 + %c0_f32x4 = vector.constant 0.0 : vector<4xf32> + %c0_offset = index.constant 0 : offset + %group_bytes = index.constant 144 : offset + %scratch_d_offset = index.constant 1024 : offset + %scratch_bytes = index.constant 1152 : offset + %token_ids_na, %weight_na, %norm_weight_na, %raw_output_na, %output_na, %q8_output_na, %f16_output_na = buffer.assume.noalias %token_ids, %weight, %norm_weight, %raw_output, %output, %q8_output, %f16_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_na[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %token_id_raw = view.load %token_ids_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %valid_token_id = scalar.cmpi sge, %token_id_raw, %c0_i32 : i32 + %safe_token_id_i32 = scf.select %valid_token_id, %token_id_raw, %c0_i32 : i32 + %token_id = index.cast %safe_token_id_i32 : i32 to index + %raw_output_view = buffer.view %raw_output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %output_view = buffer.view %output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %norm_weight_view = buffer.view %norm_weight_na[%c0_offset] : buffer -> view<[%bounded_hidden_size]xf32> + %f16_output_view = buffer.view %f16_output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf16> + %channel0 = index.mul %workitem, %c4 : index + %squares = scf.for %channel = [%channel0 to %bounded_hidden_size step %c1024](%running = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) { + %values = template.apply<@ggml.get_rows_f32.load_vector4_common>(%weight_format, %valid_token_id, %valid_token_id, %token_id, %bounded_row_count, %bounded_hidden_size, %channel, %weight_na) : (index, i1, i1, index, index, index, index, buffer) -> (vector<4xf32>) + vector.store %values, %raw_output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %square = vector.mulf %values, %values : vector<4xf32> + %next = vector.addf %running, %square : vector<4xf32> + scf.yield %next : vector<4xf32> + } + %thread_sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 + %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 + %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_offset] : buffer -> view<32xf32> + %subgroup_leader = index.cmp eq, %lane, %c0 : index + scf.if %subgroup_leader { + view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %first_subgroup = index.cmp eq, %subgroup, %c0 : index + %partial_lane = index.cmp ult, %lane, %c8 : index + %loads_partial = scalar.andi %first_subgroup, %partial_lane : i1 + %partial = scf.if %loads_partial -> (f32) { + %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 + scf.yield %value : f32 + } else { + scf.yield %c0_f32 : f32 + } + %row_sum = kernel.subgroup.reduce %partial : f32 + %writes_scale = scalar.andi %first_subgroup, %subgroup_leader : i1 + scf.if %writes_scale { + %hidden_i32 = index.cast %bounded_hidden_size : index to i32 + %hidden_f32 = scalar.sitofp %hidden_i32 : i32 to f32 + %mean = scalar.divf %row_sum, %hidden_f32 : f32 + %biased_mean = scalar.addf %mean, %epsilon : f32 + %scale = scalar.rsqrtf %biased_mean : f32 + view.store %scale, %scratch_values[%c0] : f32, view<256xf32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 + %scale_vector = vector.splat %scale : vector<4xf32> + %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 + scf.for %channel = [%channel0 to %bounded_hidden_size step %c1024] { + %values = vector.load %raw_output_view[%token, %channel] : view<[%bounded_token_count]x[%bounded_hidden_size]xf32> -> vector<4xf32> + %weights = vector.load %norm_weight_view[%channel] : view<[%bounded_hidden_size]xf32> -> vector<4xf32> + %normalized = vector.mulf %values, %scale_vector : vector<4xf32> + %result = vector.mulf %normalized, %weights : vector<4xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %half = vector.fptrunc %result : vector<4xf32> to vector<4xf16> + vector.store %half, %f16_output_view[%token, %channel] : vector<4xf16>, view<[%bounded_token_count]x[%bounded_hidden_size]xf16> + template.apply<@ggml.normalization.publish_q8.binary_vector4>(%valid_token_id, %token_output_byte_base, %channel, %result, %scratch_values, %scratch_d, %q8_output_na) : (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/rows/get_rows_scale_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scale_f32.loom new file mode 100644 index 000000000000..b229d9968b9b --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/get_rows_scale_f32.loom @@ -0,0 +1,90 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @ggml.get_rows_f32.load_vector4_common(%weight_format: index, %publish: i1, %valid_token_id: i1, %token: index, %row_count: index, %hidden_size: index, %channel: index, %weight: buffer) -> (vector<4xf32>) + +amdgpu.target @ggml_get_rows_scale_f32_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.get_rows_f32.token_capacity : %value: index where [range(%value, 1, 2048)] +config.decl @ggml.get_rows_f32.hidden_capacity : %value: index where [range(%value, 4, 1073741824), mul(%value, 4)] +config.decl @ggml.get_rows_f32.weight_format : %value: index where [range(%value, 4, 81)] +config.decl @ggml.get_rows_scale_f32.scale : f32 +config.decl @ggml.get_rows_scale_f32.bias : f32 + +kernel.def target(@ggml_get_rows_scale_f32_gfx11_wave64) export("ggml_get_rows_scale_f32") @ggml_get_rows_scale_f32(%token_count: index, %row_count: index, %hidden_size: index) { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %c256 : index + kernel.launch.config workgroups(%packet_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %raw_output: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + %token_capacity = config.get @ggml.get_rows_f32.token_capacity : index + %hidden_capacity = config.get @ggml.get_rows_f32.hidden_capacity : index + %weight_format = config.get @ggml.get_rows_f32.weight_format : index + %scale = config.get @ggml.get_rows_scale_f32.scale : f32 + %bias = config.get @ggml.get_rows_scale_f32.bias : f32 + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_row_count = index.assume %row_count [range(%row_count, 1, 262208)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 4, 1073741824), mul(%hidden_size, 4), le(%hidden_size, %hidden_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c0_offset = index.constant 0 : offset + %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index + %packet_tile_base = index.mul %channel_tile, %c256 : index + %packet0 = index.add %packet_tile_base, %workitem : index + %packet_count = index.div %bounded_hidden_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %valid_packet = index.cmp ult, %packet0, %packet_count : index + %publish = scalar.andi %valid_token, %valid_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %packet = scf.select %valid_packet, %packet0, %c0 : index + %token_ids_na, %weight_na, %raw_output_na, %output_na = buffer.assume.noalias %token_ids, %weight, %raw_output, %output : buffer, buffer, buffer, buffer + %token_ids_view = buffer.view %token_ids_na[%c0_offset] : buffer -> view<[%bounded_token_count]xi32> + %token_id_raw = view.load %token_ids_view[%token] : view<[%bounded_token_count]xi32> -> i32 + %token_id_nonnegative = scalar.cmpi sge, %token_id_raw, %c0_i32 : i32 + %safe_token_id_i32 = scf.select %token_id_nonnegative, %token_id_raw, %c0_i32 : i32 + %token_id = index.cast %safe_token_id_i32 : i32 to index + %channel0 = index.mul %packet, %c4 : index + %channel_end0 = index.add %channel0, %c4 : index + %channel, %channel_end = index.assume %channel0, %channel_end0 [le(%channel_end0, %bounded_hidden_size)] : index, index + %values = template.apply<@ggml.get_rows_f32.load_vector4_common>(%weight_format, %publish, %token_id_nonnegative, %token_id, %bounded_row_count, %bounded_hidden_size, %channel, %weight_na) : (index, i1, i1, index, index, index, index, buffer) -> (vector<4xf32>) + %scale_vector = vector.splat %scale : vector<4xf32> + %bias_vector = vector.splat %bias : vector<4xf32> + %scaled = vector.mulf %values, %scale_vector : vector<4xf32> + %result = vector.addf %scaled, %bias_vector : vector<4xf32> + scf.if %publish { + %raw_output_view = buffer.view %raw_output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + %output_view = buffer.view %output_na[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + vector.store %values, %raw_output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + vector.store %result, %output_view[%token, %channel] : vector<4xf32>, view<[%bounded_token_count]x[%bounded_hidden_size]xf32> + } + kernel.return +} + +check.case public @ggml_get_rows_scale_f32_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(4) : index + %token_ids = check.generate.iota offset(3) step(-1) : tensor<4xi32> + %weight = check.generate.fill value(2.0) : tensor<4x4xf32> + %raw_output = check.generate.fill value(0.0) : tensor<4x4xf32> + %output = check.generate.fill value(0.0) : tensor<4x4xf32> + %expected_raw = check.generate.fill value(2.0) : tensor<4x4xf32> + %expected = check.generate.fill value(4.25) : tensor<4x4xf32> + kernel.launch @ggml_get_rows_scale_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %raw_output, %output) : [index, index, index](index, index, index, tensor<4xi32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>) + check.expect.close actual(%raw_output) expected(%expected_raw) atol(0.0) rtol(0.0) nan(same) : tensor<4x4xf32> + check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<4x4xf32> + check.return +} + +check.benchmark<@ggml_get_rows_scale_f32_f32_case> @ggml_get_rows_scale_f32_f32 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom new file mode 100644 index 000000000000..46bcb16db5a8 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/rows/set_rows.loom @@ -0,0 +1,108 @@ +// Generic GGML SET_ROWS for contiguous 2D row caches. +template.decl @ggml.set_rows.body(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) + +template.decl @ggml.set_rows.launch(%hidden_capacity: index, %token_capacity: index) -> (index, index, index, index) + +template.decl @ggml.set_rows.load_f32_vector4(%input_format: index, %token: index, %channel: index, %token_count: index, %input_stride: index, %rows: buffer) -> (vector<4xf32>) + +template.decl @ggml.rows.cache_publish.safe_i64_row(%arg0: index, %arg1: index, %arg2: index, %arg3: buffer) -> (i1, index) + +template.decl @ggml.rows.cache_publish.f32_vector4(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: vector<4xf32>, %arg7: buffer) + +amdgpu.target @ggml_set_rows_gfx11_wave64 {subgroup_size = 64} + +config.decl @ggml.set_rows.token_capacity : %value: index where [range(%value, 1, 2048)] + +config.decl @ggml.set_rows.hidden_capacity : %value: index where [range(%value, 4, 32768), mul(%value, 4)] + +config.decl @ggml.set_rows.input_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.set_rows.output_format : %value: index where [range(%value, 16, 32)] + +config.decl @ggml.set_rows.input_stride : %value: index where [range(%value, 4, 1048576)] + +template.def<@ggml.set_rows.load_f32_vector4> device @ggml_set_rows_load_f32_vector4(%input_format: index, %token: index, %channel: index, %token_count: index, %input_stride: index, %rows: buffer) -> (vector<4xf32>) { + %c16 = index.constant 16 : index + %c32 = index.constant 32 : index + %c0_offset = index.constant 0 : offset + %c0_f32 = scalar.constant 0.0 : f32 + %zero = vector.splat %c0_f32 : vector<4xf32> + %is_f16 = index.cmp eq, %input_format, %c16 : index + %is_f32 = index.cmp eq, %input_format, %c32 : index + %f16_values = scf.if %is_f16 -> (vector<4xf32>) { + %rows_view = buffer.view %rows[%c0_offset] : buffer -> view<[%token_count]x[%input_stride]xf16> + %loaded = vector.load %rows_view[%token, %channel] : view<[%token_count]x[%input_stride]xf16> -> vector<4xf16> + %widened = vector.extf %loaded : vector<4xf16> to vector<4xf32> + scf.yield %widened : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %f32_values = scf.if %is_f32 -> (vector<4xf32>) { + %rows_view = buffer.view %rows[%c0_offset] : buffer -> view<[%token_count]x[%input_stride]xf32> + %loaded = vector.load %rows_view[%token, %channel] : view<[%token_count]x[%input_stride]xf32> -> vector<4xf32> + scf.yield %loaded : vector<4xf32> + } else { + scf.yield %zero : vector<4xf32> + } + %selected = scf.select %is_f32, %f32_values, %f16_values : vector<4xf32> + template.return %selected : vector<4xf32> +} + +template.def<@ggml.set_rows.body> device @ggml_set_rows_body(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) { + %token_capacity = config.get @ggml.set_rows.token_capacity : index + %hidden_capacity = config.get @ggml.set_rows.hidden_capacity : index + %input_format = config.get @ggml.set_rows.input_format : index + %output_format = config.get @ggml.set_rows.output_format : index + %input_stride = config.get @ggml.set_rows.input_stride : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index + %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index + %bounded_hidden_size = index.assume %hidden_size [range(%hidden_size, 4, 32768), mul(%hidden_size, 4), le(%hidden_size, %hidden_capacity)] : index + %channel_tile = kernel.workgroup.id : index + %token0 = kernel.workgroup.id : index + %workitem0 = kernel.workitem.id : index + %c0 = index.constant 0 : index + %c4 = index.constant 4 : index + %c256 = index.constant 256 : index + %packet_tile_base = index.mul %channel_tile, %c256 : index + %packet0 = index.add %packet_tile_base, %workitem0 : index + %packet_count = index.div %bounded_hidden_size, %c4 : index + %valid_token = index.cmp ult, %token0, %bounded_token_count : index + %valid_packet = index.cmp ult, %packet0, %packet_count : index + %publish0 = scalar.andi %valid_token, %valid_packet : i1 + %token = scf.select %valid_token, %token0, %c0 : index + %packet = scf.select %valid_packet, %packet0, %c0 : index + %channel = index.mul %packet, %c4 : index + %rows_noalias, %indices_noalias, %cache_noalias = buffer.assume.noalias %rows, %indices, %cache : buffer, buffer, buffer + %valid_row, %cache_row = template.apply<@ggml.rows.cache_publish.safe_i64_row>(%token, %bounded_token_count, %bounded_cache_row_count, %indices_noalias) : (index, index, index, buffer) -> (i1, index) + %publish = scalar.andi %publish0, %valid_row : i1 + %values = template.apply<@ggml.set_rows.load_f32_vector4>(%input_format, %token, %channel, %bounded_token_count, %input_stride, %rows_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>) + template.apply<@ggml.rows.cache_publish.f32_vector4>(%output_format, %publish, %cache_row, %channel, %bounded_cache_row_count, %bounded_hidden_size, %values, %cache_noalias) : (index, i1, index, index, index, index, vector<4xf32>, buffer) + template.return +} + +template.def<@ggml.set_rows.launch> @ggml_set_rows_launch(%hidden_capacity: index, %token_capacity: index) -> (index, index, index, index) { + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %c256 = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %c256 : index + template.return %packet_tiles, %token_capacity, %c1, %c256 : index, index, index, index +} + +kernel.def target(@ggml_set_rows_gfx11_wave64) @ggml_set_rows(%token_count: index, %cache_row_count: index, %hidden_size: index) { + %token_capacity = config.get @ggml.set_rows.token_capacity : index + %hidden_capacity = config.get @ggml.set_rows.hidden_capacity : index + %c1 = index.constant 1 : index + %c4 = index.constant 4 : index + %c255 = index.constant 255 : index + %workgroup_size = index.constant 256 : index + %packet_count = index.div %hidden_capacity, %c4 : index + %padded_packet_count = index.add %packet_count, %c255 : index + %packet_tiles = index.div %padded_packet_count, %workgroup_size : index + kernel.launch.config workgroups(%packet_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %cache_row_count: index, %hidden_size: index, %rows: buffer, %indices: buffer, %cache: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.set_rows.body>(%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : (index, index, index, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/decode_dconv4_silu_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/decode_dconv4_silu_f32.loom new file mode 100644 index 000000000000..bdabc29bfe09 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/decode_dconv4_silu_f32.loom @@ -0,0 +1,24 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.state_materialized_body(%d_inner0: index, %n_t0: index, %n_s0: index, %cache_count0: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) + +config.decl @llm.ssm_conv.decode.d_inner : %value: index where [range(%value, 32, 65536), mul(%value, 32)] + +config.decl @llm.ssm_conv.decode.workgroup_size : %value: index where [range(%value, 256, 256)] + +kernel.def export("llm_ssm_conv_dconv4_silu_decode_f32") @llm_ssm_conv_dconv4_silu_decode_f32() { + %c1 = index.constant 1 : index + %cneg1 = index.constant -1 : index + %d_inner = config.get @llm.ssm_conv.decode.d_inner : index + %wg = config.get @llm.ssm_conv.decode.workgroup_size : index + %rounding = index.add %wg, %cneg1 : index + %rounded = index.add %d_inner, %rounding : index + %groups = index.div %rounded, %wg : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%wg, %c1, %c1) : index +} launch(%state: buffer, %x: buffer, %filt: buffer, %dst: buffer, %cache: buffer) { + %c1 = index.constant 1 : index + %d_inner = config.get @llm.ssm_conv.decode.d_inner : index + template.apply<@llm.ssm_conv.state_materialized_body>(%d_inner, %c1, %c1, %c1, %state, %x, %filt, %dst, %cache, %cache, %cache, %cache, %cache) : (index, index, 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/ssm_conv/generic_binary_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_binary_f32.loom new file mode 100644 index 000000000000..aab4e830eb2e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_binary_f32.loom @@ -0,0 +1,85 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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) + +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.unary_op : %value: index where [range(%value, 0, 23)] + +config.decl @llm.ssm_conv.generic.binary_op : %value: index where [range(%value, 0, 8)] + +config.decl @llm.ssm_conv.generic.binary_lhs : %value: index where [range(%value, 0, 1)] + +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(%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 + %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 + %rounding = index.add %wg, %cneg1 : index + %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 + %c1 = index.constant 1 : index + + %d_conv0 = config.get @llm.ssm_conv.generic.d_conv : index + %d_inner0 = config.get @llm.ssm_conv.generic.d_inner : 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 + %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, 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 + %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 + %operand_g = buffer.assume.memory_space %operand : buffer + %output_g = buffer.assume.memory_space %output : buffer + %window_na, %filter_na, %operand_na, %output_na = buffer.assume.noalias %window_g, %filter_g, %operand_g, %output_g : buffer, buffer, buffer, buffer + %window_v = buffer.view %window_na[%base] : buffer -> view<1073741824xf32> + %filter_v = buffer.view %filter_na[%base] : buffer -> view<1048576xf32> + %operand_v = buffer.view %operand_na[%base] : buffer -> view<134217728xf32> + %output_v = buffer.view %output_na[%base] : buffer -> view<134217728xf32> + + scf.if %in_bounds { + %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 + %rhs = scf.select %is_lhs, %operand_value, %activated : f32 + %result = template.apply<@ggml.binary_f32.apply>(%binary_op, %lhs, %rhs) : (index, f32, f32) -> (f32) + view.store %result, %output_v[%linear] : f32, view<134217728xf32> + } + kernel.return +} 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 new file mode 100644 index 000000000000..a8ad12a57057 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/generic_f32.loom @@ -0,0 +1,69 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +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.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(%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 + %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 + %rounding = index.add %wg, %cneg1 : index + %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 + + %d_conv0 = config.get @llm.ssm_conv.generic.d_conv : index + %d_inner0 = config.get @llm.ssm_conv.generic.d_inner : 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 + %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, 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 + %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 + %output_g = buffer.assume.memory_space %output : buffer + %window_na, %filter_na, %output_na = buffer.assume.noalias %window_g, %filter_g, %output_g : buffer, buffer, buffer + %window_v = buffer.view %window_na[%base] : buffer -> view<1073741824xf32> + %filter_v = buffer.view %filter_na[%base] : buffer -> view<1048576xf32> + %output_v = buffer.view %output_na[%base] : buffer -> view<134217728xf32> + + scf.if %in_bounds { + %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/loom-libs/ops/ssm_conv/lfm_dconv3_state_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/lfm_dconv3_state_f32.loom new file mode 100644 index 000000000000..a0ffcb1a08cd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/lfm_dconv3_state_f32.loom @@ -0,0 +1,128 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @llm_ssm_conv_lfm_dconv3_gfx11_wave32 {subgroup_size = 32} + +config.decl @llm.ssm_conv.lfm_dconv3.n_t : %value: index where [range(%value, 1, 64)] + +kernel.def target(@llm_ssm_conv_lfm_dconv3_gfx11_wave32) export("llm_ssm_conv_lfm_dconv3_state_f32") @llm_ssm_conv_lfm_dconv3_state_f32() { + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %n_t = config.get @llm.ssm_conv.lfm_dconv3.n_t : index + %groups = index.mul %n_t, %c8 : index + kernel.launch.config workgroups(%groups, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%state: buffer, %x: buffer, %filter: buffer, %output: buffer, %cache: buffer) { + %base = index.constant 0 : offset + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c8 = index.constant 8 : index + %c2048 = index.constant 2048 : index + %c256 = index.constant 256 : index + %n_t0 = config.get @llm.ssm_conv.lfm_dconv3.n_t : index + %n_t = index.assume %n_t0 [range(%n_t0, 1, 64)] : index + %group0 = kernel.workgroup.id : index + %lane0 = kernel.workitem.id : index + %group = index.assume %group0 [range(%group0, 0, 511)] : index + %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index + %token0 = index.div %group, %c8 : index + %token = index.assume %token0 [range(%token0, 0, 63)] : index + %channel_group = index.rem %group, %c8 : index + %channel_base = index.mul %channel_group, %c256 : index + %channel0 = index.add %channel_base, %lane : index + %channel = index.assume %channel0 [range(%channel0, 0, 2047)] : index + %token_elements0 = index.mul %n_t, %c2048 : index + %token_elements = index.assume %token_elements0 [range(%token_elements0, 2048, 131072), mul(%token_elements0, 2048)] : index + + %state_v = buffer.view %state[%base] : buffer -> view<4096xf32> + %filter_v = buffer.view %filter[%base] : buffer -> view<6144xf32> + %cache_v = buffer.view %cache[%base] : buffer -> view<4096xf32> + %state_base0 = index.mul %channel, %c2 : index + %state_base = index.assume %state_base0 [range(%state_base0, 0, 4094), mul(%state_base0, 2)] : index + %filter_base0 = index.mul %channel, %c3 : index + %filter_base = index.assume %filter_base0 [range(%filter_base0, 0, 6141), mul(%filter_base0, 3)] : index + %filter_one0 = index.add %filter_base, %c1 : index + %filter_one = index.assume %filter_one0 [range(%filter_one0, 1, 6142)] : index + %filter_two0 = index.add %filter_base, %c2 : index + %filter_two = index.assume %filter_two0 [range(%filter_two0, 2, 6143)] : index + %w0 = view.load %filter_v[%filter_base] : view<6144xf32> -> f32 + %w1 = view.load %filter_v[%filter_one] : view<6144xf32> -> f32 + %w2 = view.load %filter_v[%filter_two] : view<6144xf32> -> f32 + + %from_x0 = index.cmp uge, %token, %c2 : index + %v0 = scf.if %from_x0 -> (f32) { + %x_token = index.sub %token, %c2 : index + %x_base = index.mul %x_token, %c2048 : index + %x_index0 = index.add %x_base, %channel : index + %x_index, %x_extent = index.assume %x_index0, %token_elements [lt(%x_index0, %token_elements)] : index, index + %x_view = buffer.view %x[%base] : buffer -> view<[%x_extent]xf32> + %loaded = view.load %x_view[%x_index] : view<[%x_extent]xf32> -> f32 + scf.yield %loaded : f32 + } else { + %state_index0 = index.add %state_base, %token : index + %state_index = index.assume %state_index0 [range(%state_index0, 0, 4095)] : index + %loaded = view.load %state_v[%state_index] : view<4096xf32> -> f32 + scf.yield %loaded : f32 + } + %position1 = index.add %token, %c1 : index + %from_x1 = index.cmp uge, %position1, %c2 : index + %v1 = scf.if %from_x1 -> (f32) { + %x_token = index.sub %position1, %c2 : index + %x_base = index.mul %x_token, %c2048 : index + %x_index0 = index.add %x_base, %channel : index + %x_index, %x_extent = index.assume %x_index0, %token_elements [lt(%x_index0, %token_elements)] : index, index + %x_view = buffer.view %x[%base] : buffer -> view<[%x_extent]xf32> + %loaded = view.load %x_view[%x_index] : view<[%x_extent]xf32> -> f32 + scf.yield %loaded : f32 + } else { + %state_index0 = index.add %state_base, %position1 : index + %state_index = index.assume %state_index0 [range(%state_index0, 1, 4095)] : index + %loaded = view.load %state_v[%state_index] : view<4096xf32> -> f32 + scf.yield %loaded : f32 + } + %x_base2 = index.mul %token, %c2048 : index + %x_index20 = index.add %x_base2, %channel : index + %x_index2, %x_extent2 = index.assume %x_index20, %token_elements [lt(%x_index20, %token_elements)] : index, index + %x_view2 = buffer.view %x[%base] : buffer -> view<[%x_extent2]xf32> + %v2 = view.load %x_view2[%x_index2] : view<[%x_extent2]xf32> -> f32 + %product0 = scalar.mulf %v0, %w0 : f32 + %product1 = scalar.mulf %v1, %w1 : f32 + %sum0 = scalar.addf %product0, %product1 : f32 + %product2 = scalar.mulf %v2, %w2 : f32 + %sum1 = scalar.addf %sum0, %product2 : f32 + %output_index, %output_extent = index.assume %x_index20, %token_elements [lt(%x_index20, %token_elements)] : index, index + %output_view = buffer.view %output[%base] : buffer -> view<[%output_extent]xf32> + view.store %sum1, %output_view[%output_index] : f32, view<[%output_extent]xf32> + + %publish = index.cmp eq, %token, %c0 : index + scf.if %publish { + %one_token = index.cmp eq, %n_t, %c1 : index + %tail0, %tail1 = scf.if %one_token -> (f32, f32) { + %state_one0 = index.add %state_base, %c1 : index + %state_one = index.assume %state_one0 [range(%state_one0, 1, 4095)] : index + %old1 = view.load %state_v[%state_one] : view<4096xf32> -> f32 + scf.yield %old1, %v2 : f32, f32 + } else { + %tail_token0 = index.sub %n_t, %c2 : index + %tail_token1 = index.sub %n_t, %c1 : index + %tail_base0 = index.mul %tail_token0, %c2048 : index + %tail_base1 = index.mul %tail_token1, %c2048 : index + %tail_index00 = index.add %tail_base0, %channel : index + %tail_index10 = index.add %tail_base1, %channel : index + %tail_index0, %tail_extent0 = index.assume %tail_index00, %token_elements [lt(%tail_index00, %token_elements)] : index, index + %tail_index1, %tail_extent1 = index.assume %tail_index10, %token_elements [lt(%tail_index10, %token_elements)] : index, index + %tail_view0 = buffer.view %x[%base] : buffer -> view<[%tail_extent0]xf32> + %tail_view1 = buffer.view %x[%base] : buffer -> view<[%tail_extent1]xf32> + %loaded0 = view.load %tail_view0[%tail_index0] : view<[%tail_extent0]xf32> -> f32 + %loaded1 = view.load %tail_view1[%tail_index1] : view<[%tail_extent1]xf32> -> f32 + scf.yield %loaded0, %loaded1 : f32, f32 + } + %cache_one0 = index.add %state_base, %c1 : index + %cache_one = index.assume %cache_one0 [range(%cache_one0, 1, 4095)] : index + view.store %tail0, %cache_v[%state_base] : f32, view<4096xf32> + view.store %tail1, %cache_v[%cache_one] : f32, view<4096xf32> + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/prefill_dconv4_silu_512_wg1024.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/prefill_dconv4_silu_512_wg1024.loom new file mode 100644 index 000000000000..51ef92db20fd --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/prefill_dconv4_silu_512_wg1024.loom @@ -0,0 +1,44 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.prefill_window(%packet_width: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer) + +config.decl @llm.ssm_conv.prefill.d_conv : %value: index where [range(%value, 4, 4)] + +config.decl @llm.ssm_conv.prefill.d_inner : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +config.decl @llm.ssm_conv.prefill.n_t : %value: index where [range(%value, 512, 512)] + +config.decl @llm.ssm_conv.prefill.n_s : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.prefill.state_row_stride : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +config.decl @llm.ssm_conv.prefill.x_row_stride : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +func.def inline @llm_ssm_conv_prefill_channel_packet(%channels: index) -> (index) { + %zero = index.constant 0 : index + %one = index.constant 1 : index + %two = index.constant 2 : index + %alignment = index.constant 64 : index + %remainder = index.rem %channels, %alignment : index + %aligned = index.cmp eq, %remainder, %zero : index + %width = scf.select %aligned, %two, %one : index + func.return %width : index +} + +kernel.def export("llm_ssm_conv_dconv4_silu_prefill_512_wg1024") @llm_ssm_conv_dconv4_silu_prefill_512_wg1024() { + %c1 = index.constant 1 : index + %channel_lanes = index.constant 32 : index + %threads = index.constant 1024 : index + %d_inner = config.get @llm.ssm_conv.prefill.d_inner : index + %n_s = config.get @llm.ssm_conv.prefill.n_s : index + %packet_width = func.call pure @llm_ssm_conv_prefill_channel_packet(%d_inner) : (index) -> (index) + %channels = index.mul %packet_width, %channel_lanes : index + %groups_x = index.div %d_inner, %channels : index + kernel.launch.config workgroups(%groups_x, %n_s, %c1) workgroup_size(%threads, %c1, %c1) : index +} launch(%state: buffer, %x: buffer, %filt: buffer, %dst: buffer) { + %d_inner = config.get @llm.ssm_conv.prefill.d_inner : index + %packet_width = func.call pure @llm_ssm_conv_prefill_channel_packet(%d_inner) : (index) -> (index) + template.apply<@llm.ssm_conv.prefill_window>(%packet_width, %state, %x, %filt, %dst) : (index, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/prefill_finish_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/prefill_finish_f32.loom new file mode 100644 index 000000000000..205bc9433518 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/prefill_finish_f32.loom @@ -0,0 +1,21 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.prefill_finish(%state: buffer, %filter: buffer, %edges: buffer, %output: buffer, %cache: buffer) + +config.decl @llm.ssm_conv.generic.d_inner : %value: index where [range(%value, 32, 65536), mul(%value, 32)] + +amdgpu.target @llm_ssm_conv_prefill_finish_wave32 {subgroup_size = 32} + +kernel.def target(@llm_ssm_conv_prefill_finish_wave32) @llm_ssm_conv_dconv4_silu_prefill_finish_f32() { + %width = config.get @llm.ssm_conv.generic.d_inner : index + %one = index.constant 1 : index + %threads = index.constant 256 : index + %round = index.constant 255 : index + %rounded = index.add %width, %round : index + %groups = index.div %rounded, %threads : index + kernel.launch.config workgroups(%groups, %one, %one) workgroup_size(%threads, %one, %one) : index +} launch(%state: buffer, %filter: buffer, %edges: buffer, %output: buffer, %cache: buffer) { + template.apply<@llm.ssm_conv.prefill_finish>(%state, %filter, %edges, %output, %cache) : (buffer, buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/rollback_dconv4_silu_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/rollback_dconv4_silu_f32.loom new file mode 100644 index 000000000000..888bb046f360 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/rollback_dconv4_silu_f32.loom @@ -0,0 +1,33 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.state_materialized_body(%d_inner0: index, %n_t0: index, %n_s0: index, %cache_count0: index, %state: buffer, %x: buffer, %filt: buffer, %dst: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) + +config.decl @llm.ssm_conv.rollback.d_inner : %value: index where [range(%value, 32, 65536), mul(%value, 32)] + +config.decl @llm.ssm_conv.rollback.n_t : %value: index where [range(%value, 1, 512)] + +config.decl @llm.ssm_conv.rollback.n_s : %value: index where [range(%value, 1, 3)] + +config.decl @llm.ssm_conv.rollback.cache_count : %value: index where [range(%value, 1, 5)] + +config.decl @llm.ssm_conv.rollback.workgroup_size : %value: index where [range(%value, 256, 256)] + +kernel.def export("llm_ssm_conv_dconv4_silu_rollback_f32") @llm_ssm_conv_dconv4_silu_rollback_f32() { + %c1 = index.constant 1 : index + %cneg1 = index.constant -1 : index + %d_inner = config.get @llm.ssm_conv.rollback.d_inner : index + %n_s = config.get @llm.ssm_conv.rollback.n_s : index + %wg = config.get @llm.ssm_conv.rollback.workgroup_size : index + %rounding = index.add %wg, %cneg1 : index + %rounded = index.add %d_inner, %rounding : index + %groups = index.div %rounded, %wg : index + kernel.launch.config workgroups(%groups, %n_s, %c1) workgroup_size(%wg, %c1, %c1) : index +} launch(%state: buffer, %x: buffer, %filt: buffer, %dst: buffer, %cache0: buffer, %cache1: buffer, %cache2: buffer, %cache3: buffer, %cache4: buffer) { + %d_inner = config.get @llm.ssm_conv.rollback.d_inner : index + %n_t = config.get @llm.ssm_conv.rollback.n_t : index + %n_s = config.get @llm.ssm_conv.rollback.n_s : index + %cache_count = config.get @llm.ssm_conv.rollback.cache_count : index + template.apply<@llm.ssm_conv.state_materialized_body>(%d_inner, %n_t, %n_s, %cache_count, %state, %x, %filt, %dst, %cache0, %cache1, %cache2, %cache3, %cache4) : (index, index, 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/ssm_conv/snapshot_window_tail_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/snapshot_window_tail_f32.loom new file mode 100644 index 000000000000..a50ed243c848 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/ssm_conv/snapshot_window_tail_f32.loom @@ -0,0 +1,47 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @llm.ssm_conv.snapshot_window_tail(%state: buffer, %x: buffer, %dst: buffer, %cache: buffer) + +config.decl @llm.ssm_conv.snapshot.d_conv : %value: index where [range(%value, 4, 4)] + +config.decl @llm.ssm_conv.snapshot.d_inner : %value: index where [range(%value, 8192, 10240), mul(%value, 32)] + +config.decl @llm.ssm_conv.snapshot.n_t : %value: index where [range(%value, 512, 512)] + +config.decl @llm.ssm_conv.snapshot.n_s : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.snapshot.state_row_stride : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.snapshot.state_channel_stride : %value: index where [range(%value, 3, 3)] + +config.decl @llm.ssm_conv.snapshot.x_row_stride : %value: index where [range(%value, 1, 16777216)] + +config.decl @llm.ssm_conv.snapshot.dst_row_stride : %value: index where [range(%value, 1, 16777216)] + +config.decl @llm.ssm_conv.snapshot.cache_row_stride : %value: index where [range(%value, 1, 1)] + +config.decl @llm.ssm_conv.snapshot.cache_channel_stride : %value: index where [range(%value, 3, 3)] + +config.decl @llm.ssm_conv.snapshot.workgroup_size : %value: index where [range(%value, 32, 1024), mul(%value, 32)] + +kernel.def export("llm_ssm_conv_snapshot_window_tail_f32") @llm_ssm_conv_snapshot_window_tail_f32() { + %unit = index.constant 1 : index + %cneg1 = index.constant -1 : index + %x_snapshot_rows = index.constant 61 : index + %d_conv = config.get @llm.ssm_conv.snapshot.d_conv : index + %d_inner = config.get @llm.ssm_conv.snapshot.d_inner : index + %n_s = config.get @llm.ssm_conv.snapshot.n_s : index + %wg = config.get @llm.ssm_conv.snapshot.workgroup_size : index + %carry = index.add %d_conv, %cneg1 : index + %snapshot_rows = index.add %carry, %x_snapshot_rows : index + %rows = index.mul %snapshot_rows, %n_s : index + %total = index.mul %rows, %d_inner : index + %rounding = index.add %wg, %cneg1 : index + %rounded = index.add %total, %rounding : index + %groups = index.div %rounded, %wg : index + kernel.launch.config workgroups(%groups, %unit, %unit) workgroup_size(%wg, %unit, %unit) : index +} launch(%state: buffer, %x: buffer, %dst: buffer, %cache: buffer) { + template.apply<@llm.ssm_conv.snapshot_window_tail>(%state, %x, %dst, %cache) : (buffer, buffer, buffer, buffer) + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom similarity index 98% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom index a75a3fa0059d..4eebda0ccda2 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata.loom @@ -78,7 +78,7 @@ kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_attention_ // Publishes the single position and cache row needed by one decode issue. // Decode attention bounds itself with the same request control word and does // not need a materialized causal mask: every prior row is visible. -kernel.def target(@qwen_attention_metadata_gfx11_wave64) export("qwen_decode_attention_metadata") @qwen_decode_attention_metadata() { +kernel.def target(@qwen_attention_metadata_gfx11_wave64) @qwen_decode_attention_metadata() { %c1 = index.constant 1 : index kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index } launch(%control: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer) { @@ -182,3 +182,5 @@ check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadat check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_128 {token_count = 128} check.benchmark<@qwen_attention_metadata_benchmark_case> @qwen_attention_metadata_prefill_512 {token_count = 512} + +check.benchmark<@qwen_decode_attention_metadata_sequence_case> @qwen_decode_attention_metadata_sequence diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata_bringup_workaround.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata_bringup_workaround.loom similarity index 87% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata_bringup_workaround.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata_bringup_workaround.loom index 079c1465fefb..6c06d4baa682 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_metadata_bringup_workaround.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_metadata_bringup_workaround.loom @@ -94,7 +94,7 @@ check.case public @qwen_attention_metadata_causal_and_nonzero_base_case { %expected_zero = check.generate.fill value(0) : tensor<1xi32> %expected_zero_indices = check.generate.fill value(0) : tensor<1xi64> %expected_zero_mask = check.generate.iota offset(0) step(-1024) : tensor<2xi16> - func.call @qwen_attention_metadata_bringup_workaround(%one, %two, %zero_control, %zero_positions, %zero_key_indices, %zero_value_indices, %zero_mask) : (index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<2xi16>) + kernel.launch @qwen_attention_metadata_bringup_workaround[%one, %two](%one, %two, %zero_control, %zero_positions, %zero_key_indices, %zero_value_indices, %zero_mask) : [index, index](index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<2xi16>) check.expect.equal actual(%zero_positions) expected(%expected_zero) : tensor<1xi32> check.expect.equal actual(%zero_key_indices) expected(%expected_zero_indices) : tensor<1xi64> check.expect.equal actual(%zero_value_indices) expected(%expected_zero_indices) : tensor<1xi64> @@ -107,7 +107,7 @@ check.case public @qwen_attention_metadata_causal_and_nonzero_base_case { %expected_nonzero = check.generate.fill value(4) : tensor<1xi32> %expected_nonzero_indices = check.generate.fill value(4) : tensor<1xi64> %expected_nonzero_mask = check.generate.fill value(0) : tensor<5xi16> - func.call @qwen_attention_metadata_bringup_workaround(%one, %five, %nonzero_control, %nonzero_positions, %nonzero_key_indices, %nonzero_value_indices, %nonzero_mask) : (index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<5xi16>) + kernel.launch @qwen_attention_metadata_bringup_workaround[%one, %five](%one, %five, %nonzero_control, %nonzero_positions, %nonzero_key_indices, %nonzero_value_indices, %nonzero_mask) : [index, index](index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<5xi16>) check.expect.equal actual(%nonzero_positions) expected(%expected_nonzero) : tensor<1xi32> check.expect.equal actual(%nonzero_key_indices) expected(%expected_nonzero_indices) : tensor<1xi64> check.expect.equal actual(%nonzero_value_indices) expected(%expected_nonzero_indices) : tensor<1xi64> @@ -124,7 +124,7 @@ check.case public @qwen_attention_metadata_benchmark_case { %attention_mask = check.generate.fill value(1) : tensor<[%token_count]x[%token_count]xi16> %expected_positions = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi32> %expected_cache_indices = check.generate.iota offset(0) step(1) : tensor<[%token_count]xi64> - func.call @qwen_attention_metadata_bringup_workaround(%token_count, %token_count, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : (index, index, tensor<1xi32>, tensor<[%token_count]xi32>, tensor<[%token_count]xi64>, tensor<[%token_count]xi64>, tensor<[%token_count]x[%token_count]xi16>) + kernel.launch @qwen_attention_metadata_bringup_workaround[%token_count, %token_count](%token_count, %token_count, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<1xi32>, tensor<[%token_count]xi32>, tensor<[%token_count]xi64>, tensor<[%token_count]xi64>, tensor<[%token_count]x[%token_count]xi16>) check.expect.equal actual(%positions) expected(%expected_positions) : tensor<[%token_count]xi32> check.expect.equal actual(%key_cache_indices) expected(%expected_cache_indices) : tensor<[%token_count]xi64> check.expect.equal actual(%value_cache_indices) expected(%expected_cache_indices) : tensor<[%token_count]xi64> @@ -141,7 +141,7 @@ check.case public @qwen_attention_metadata_decode_768_case { %key_cache_indices = check.generate.fill value(-1) : tensor<1xi64> %value_cache_indices = check.generate.fill value(-1) : tensor<1xi64> %attention_mask = check.generate.fill value(1) : tensor<768xi16> - func.call @qwen_attention_metadata_bringup_workaround(%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : (index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<768xi16>) + kernel.launch @qwen_attention_metadata_bringup_workaround[%token_count, %context_capacity](%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<1xi32>, tensor<1xi32>, tensor<1xi64>, tensor<1xi64>, tensor<768xi16>) check.return } @@ -153,7 +153,7 @@ check.case public @qwen_attention_metadata_prefill_512_case { %key_cache_indices = check.generate.fill value(-1) : tensor<512xi64> %value_cache_indices = check.generate.fill value(-1) : tensor<512xi64> %attention_mask = check.generate.fill value(1) : tensor<512x512xi16> - func.call @qwen_attention_metadata_bringup_workaround(%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : (index, index, tensor<1xi32>, tensor<512xi32>, tensor<512xi64>, tensor<512xi64>, tensor<512x512xi16>) + kernel.launch @qwen_attention_metadata_bringup_workaround[%token_count, %context_capacity](%token_count, %context_capacity, %control, %positions, %key_cache_indices, %value_cache_indices, %attention_mask) : [index, index](index, index, tensor<1xi32>, tensor<512xi32>, tensor<512xi64>, tensor<512xi64>, tensor<512x512xi16>) check.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_state_initialize.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_state_initialize.loom similarity index 98% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_state_initialize.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_state_initialize.loom index bade957f43d0..9754207c7210 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/attention_state_initialize.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/attention_state_initialize.loom @@ -83,7 +83,7 @@ check.case public @qwen_attention_context_base_capture_case { %positions = check.generate.fill value(7) : tensor<1xi32> %control = check.generate.fill value(-1) : tensor<1xi32> %expected = check.generate.fill value(7) : tensor<1xi32> - kernel.launch @qwen_attention_context_base_capture[](%positions, %control) : [](tensor<1xi32>, tensor<1xi32>) + kernel.launch @qwen_attention_context_base_capture(%positions, %control) : (tensor<1xi32>, tensor<1xi32>) check.expect.equal actual(%control) expected(%expected) : tensor<1xi32> check.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/manifest.json b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/manifest.json new file mode 100644 index 000000000000..da6f8462d037 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/manifest.json @@ -0,0 +1,195 @@ +{ + "schema": "ggml-hrx-kernel-corpus-v1", + "upstream_revision": "local", + "files": [ + { + "path": "token_embedding_bringup_workaround.loom" + }, + { + "path": "attention_state_initialize.loom" + }, + { + "path": "attention_metadata_bringup_workaround.loom" + } + ], + "exports": [ + { + "name": "qwen_attention_context_base_capture", + "symbol": "qwen_attention_context_base_capture", + "source": "attention_state_initialize.loom", + "workload_parameters": [], + "launch_parameters": [], + "bindings": [ + "positions", + "control" + ], + "binding_access": [ + "read", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "attention_state_initialize.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [], + "family": "qwen" + }, + { + "name": "qwen_attention_decode_state_initialize", + "symbol": "qwen_attention_decode_state_initialize", + "source": "attention_state_initialize.loom", + "workload_parameters": [ + { + "name": "context_capacity", + "type": "index" + }, + { + "name": "completion_counter_count", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "context_capacity", + "type": "index" + }, + { + "name": "completion_counter_count", + "type": "index" + } + ], + "bindings": [ + "positions", + "key_cache_indices", + "value_cache_indices", + "attention_mask", + "completion_counters" + ], + "binding_access": [ + "read", + "write", + "write", + "write", + "write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "attention_state_initialize.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [], + "family": "qwen" + }, + { + "name": "qwen_attention_metadata_bringup_workaround", + "symbol": "qwen_attention_metadata_bringup_workaround", + "source": "attention_metadata_bringup_workaround.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "context_capacity", + "type": "index" + } + ], + "bindings": [ + "control", + "positions", + "key_cache_indices", + "value_cache_indices", + "attention_mask" + ], + "binding_access": [ + "read", + "read_write", + "read_write", + "read_write", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "attention_metadata_bringup_workaround.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [], + "family": "qwen" + }, + { + "name": "qwen_token_embedding_q4k_bringup_workaround", + "symbol": "qwen_token_embedding_q4k_bringup_workaround", + "source": "token_embedding_bringup_workaround.loom", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "vocabulary_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "launch_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "vocabulary_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ], + "bindings": [ + "token_ids", + "weight", + "output" + ], + "binding_access": [ + "read", + "read", + "read_write" + ], + "target_selector": "", + "compile_recipe": { + "mode": "direct", + "primary_sources": [ + "token_embedding_bringup_workaround.loom" + ], + "library_sources": [] + }, + "compile_dependencies": [], + "family": "qwen" + } + ], + "link_modules": [], + "plan_cases": [] +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_bringup_workaround.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/token_embedding_bringup_workaround.loom similarity index 92% rename from ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_bringup_workaround.loom rename to ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/token_embedding_bringup_workaround.loom index 8977fee1060c..2a4db9725a33 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_bringup_workaround.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen/token_embedding_bringup_workaround.loom @@ -160,7 +160,7 @@ check.case public @qwen_token_embedding_q4k_decode_case { %weight = check.generate.fill value(85) : tensor<1x8x144xi8> %output = check.generate.fill value(0.0) : tensor<1x2048xf32> %expected = check.generate.fill value(7166.25) : tensor<1x2048xf32> - func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<1xi32>, tensor<1x8x144xi8>, tensor<1x2048xf32>) + kernel.launch @qwen_token_embedding_q4k_bringup_workaround[%token_count, %vocabulary_count, %hidden_size](%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<1x8x144xi8>, tensor<1x2048xf32>) check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> check.return } @@ -175,7 +175,7 @@ check.case public @qwen_token_embedding_q4k_row_access_case { %weight = check.generate.fill value(0) : tensor<3x8x144xi8> %output = check.generate.fill value(1.0) : tensor<3x2048xf32> %expected = check.generate.fill value(0.0) : tensor<3x2048xf32> - func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<3xi32>, tensor<3x8x144xi8>, tensor<3x2048xf32>) + kernel.launch @qwen_token_embedding_q4k_bringup_workaround[%token_count, %vocabulary_count, %hidden_size](%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<3xi32>, tensor<3x8x144xi8>, tensor<3x2048xf32>) check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<3x2048xf32> check.return } @@ -188,7 +188,7 @@ check.case public @qwen_token_embedding_q4k_benchmark_case { %weight = check.generate.fill value(0) : tensor<151936x8x144xi8> %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> %expected = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<[%token_count]xi32>, tensor<151936x8x144xi8>, tensor<[%token_count]x2048xf32>) + kernel.launch @qwen_token_embedding_q4k_bringup_workaround[%token_count, %vocabulary_count, %hidden_size](%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]xi32>, tensor<151936x8x144xi8>, tensor<[%token_count]x2048xf32>) check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> check.return } @@ -202,7 +202,7 @@ check.case public @qwen35_token_embedding_q4k_row_access_case { %weight = check.generate.fill value(0) : tensor<248320x12x144xi8> %output = check.generate.fill value(1.0) : tensor<1x3072xf32> %expected = check.generate.fill value(0.0) : tensor<1x3072xf32> - func.call @qwen_token_embedding_q4k_bringup_workaround(%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : (index, index, index, tensor<1xi32>, tensor<248320x12x144xi8>, tensor<1x3072xf32>) + kernel.launch @qwen_token_embedding_q4k_bringup_workaround[%token_count, %vocabulary_count, %hidden_size](%token_count, %vocabulary_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<248320x12x144xi8>, tensor<1x3072xf32>) check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x3072xf32> check.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_f32.loom deleted file mode 100644 index af203231251c..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_f32.loom +++ /dev/null @@ -1,423 +0,0 @@ -// Contracts GGML Q6_K weight rows directly with F32 activation rows using -// the decode schedule selected by llama.cpp's Vulkan backend. -// -// One 64-workitem workgroup computes two adjacent output rows. Four cohorts -// of 16 lanes each process four Q6_K blocks in parallel, while each lane -// decodes four values from the 0, 32, 64, and 96 element quarters of its -// block. Per-group Q6_K scales are exchanged through two 256-byte LDS frames, -// one for each output row. The frames preserve the producer/consumer shape of -// the Vulkan oracle without repacking the persistent GGUF weights. -// -// This provider is intentionally independent of the Q8_1 integer-dot path. -// On gfx11, direct F32 loads and software Q6_K decode can beat activation -// quantization for decode-sized batches. Keeping both algorithms in the same -// corpus lets JIT selection depend on the specialized shape and target. -amdgpu.target @ggml_q6k_gfx11_wave64 {subgroup_size = 64} - -amdgpu.target @ggml_q6k_gfx11_wave32 {subgroup_size = 32} - -config.decl @ggml.linear_q6k_f32.token_capacity : %value: index where [range(%value, 1, 2048)] - -config.decl @ggml.linear_q6k_f32.output_capacity : %value: index where [range(%value, 1, 262144)] - -// Q8_1 reference declarations used only by the differential case. -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) - -// Computes an ordered four-element F32 dot product. The scalar recurrence -// mirrors the GLSL oracle and gives target scheduling four independent dot -// chains per decoded Q6_K block. -func.def inline @ggml_q6k_dot4_f32(%lhs: vector<4xf32>, %rhs: vector<4xf32>) -> (f32) { - %c0 = scalar.constant 0.0 : f32 - %lhs0 = vector.extract %lhs[0] : vector<4xf32> -> f32 - %lhs1 = vector.extract %lhs[1] : vector<4xf32> -> f32 - %lhs2 = vector.extract %lhs[2] : vector<4xf32> -> f32 - %lhs3 = vector.extract %lhs[3] : vector<4xf32> -> f32 - %rhs0 = vector.extract %rhs[0] : vector<4xf32> -> f32 - %rhs1 = vector.extract %rhs[1] : vector<4xf32> -> f32 - %rhs2 = vector.extract %rhs[2] : vector<4xf32> -> f32 - %rhs3 = vector.extract %rhs[3] : vector<4xf32> -> f32 - %sum0 = scalar.fmaf %lhs0, %rhs0, %c0 : f32 - %sum1 = scalar.fmaf %lhs1, %rhs1, %sum0 : f32 - %sum2 = scalar.fmaf %lhs2, %rhs2, %sum1 : f32 - %sum3 = scalar.fmaf %lhs3, %rhs3, %sum2 : f32 - func.return %sum3 : f32 -} - -// Loads one row contraction's activation packet. The two output rows read the -// same packet, but each load remains adjacent to its consumer. Hoisting the -// packet across the intervening LDS barrier halves issued activation traffic -// while extending sixteen F32 live values across synchronization; on gfx1151 -// that longer live range is slower than the repeated cache-hot read. -func.def inline @ggml_q6k_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 - %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 - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index - %block_count = index.div %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 - %input0 = vector.load %input_view[%token, %safe_block, %input_index0] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> - %input1 = vector.load %input_view[%token, %safe_block, %input_index1] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> - %input2 = vector.load %input_view[%token, %safe_block, %input_index2] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> - %input3 = vector.load %input_view[%token, %safe_block, %input_index3] : view<[%token_count]x[%launch_block_count]x256xf32> -> vector<4xf32> - func.return %input0, %input1, %input2, %input3 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32> -} - -// Loads the signed group scale owned by one lane of a Q6_K block. -func.def inline @ggml_q6k_load_f32_scale(%input_size: index, %row: index, %block: index, %lane: index, %weight: buffer) -> (f32) { - %c0 = index.constant 0 : index - %c16 = index.constant 16 : index - %c192 = index.constant 192 : offset - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index - %block_count = index.div %input_size, %c256 : index - %valid_block = index.cmp ult, %block, %block_count : index - %safe_block = scf.if %valid_block -> (index) { - scf.yield %block : index - } else { - scf.yield %c0 : index - } - %itid = index.rem %bounded_lane, %c16 : index - %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset - %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset - %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset - %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %weight_block_byte_base, %c192 : offset - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %scale_i8 = view.load %scale_view[%itid] : view<16xi8> -> i8 - %scale = scalar.sitofp %scale_i8 : i8 to f32 - func.return %scale : f32 -} - -// Stages one row's group scales and synchronizes before contraction. -func.def inline @ggml_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) { - %c0_offset = index.constant 0 : offset - %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index - %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> - %scale = func.call @ggml_q6k_load_f32_scale(%input_size, %row, %block, %bounded_lane, %weight) : (index, index, index, index, buffer) -> (f32) - view.store %scale, %scale_stage_view[%frame, %bounded_lane] : f32, view<[%frame_count]x64xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - func.return -} - -// Computes one lane contribution for one output row and one four-block -// cohort after its group scales have been staged. -func.def inline @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) { - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c6 = index.constant 6 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c128 = index.constant 128 : offset - %c208 = index.constant 208 : offset - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_i32v = vector.constant 4 : vector<1xi32> - %c2_i32v = vector.constant 2 : vector<1xi32> - %nibble_mask = vector.constant 252645135 : vector<1xi32> - %high0_mask = vector.constant 50529027 : vector<1xi32> - %high2_mask = vector.constant 202116108 : vector<1xi32> - %high4_mask = vector.constant 808464432 : vector<1xi32> - %high6_mask = vector.constant -1061109568 : vector<1xi32> - %c32_f32v = vector.constant 32.0 : vector<4xf32> - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index - %block_count = index.div %input_size, %c256 : index - %valid_block = index.cmp ult, %block, %block_count : index - %safe_block = scf.if %valid_block -> (index) { - scf.yield %block : index - } else { - scf.yield %c0 : index - } - %itid = index.rem %bounded_lane, %c16 : index - %cohort = index.div %bounded_lane, %c16 : index - %vector_half = index.div %itid, %c8 : index - %vector_index = index.rem %itid, %c8 : index - %vector_quarter = index.div %vector_index, %c4 : index - %weight_row_bytes = index.scale %block_count, %c210_bytes : index, offset -> offset - %weight_row_byte_base = index.scale %row, %weight_row_bytes : index, offset -> offset - %block_byte_add = index.scale %safe_block, %c210_bytes : index, offset -> offset - %weight_block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %weight_block_byte_base, %c128 : offset - %d_byte_base = index.add %weight_block_byte_base, %c208 : offset - %ql_view = buffer.view %weight[%weight_block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<[%frame_count]x64xf32> - %scale_cohort_base = index.mul %cohort, %c16 : index - %scale_half_base = index.mul %vector_half, %c8 : index - %scale_index0 = index.add %scale_half_base, %vector_quarter : index - %scale_index1 = index.add %scale_index0, %c2 : index - %scale_index2 = index.add %scale_index0, %c4 : index - %scale_index3 = index.add %scale_index0, %c6 : index - %stage_scale_index0 = index.add %scale_cohort_base, %scale_index0 : index - %stage_scale_index1 = index.add %scale_cohort_base, %scale_index1 : index - %stage_scale_index2 = index.add %scale_cohort_base, %scale_index2 : index - %stage_scale_index3 = index.add %scale_cohort_base, %scale_index3 : index - %scale0 = view.load %scale_stage_view[%frame, %stage_scale_index0] : view<[%frame_count]x64xf32> -> f32 - %scale1 = view.load %scale_stage_view[%frame, %stage_scale_index1] : view<[%frame_count]x64xf32> -> f32 - %scale2 = view.load %scale_stage_view[%frame, %stage_scale_index2] : view<[%frame_count]x64xf32> -> f32 - %scale3 = view.load %scale_stage_view[%frame, %stage_scale_index3] : view<[%frame_count]x64xf32> -> f32 - %ql_half_word_base = index.mul %vector_half, %c16 : index - %ql_word_index00 = index.add %ql_half_word_base, %vector_index : index - %ql_word_index10 = index.add %ql_word_index00, %c8 : index - %qh_half_word_base = index.mul %vector_half, %c8 : index - %qh_word_index0 = index.add %qh_half_word_base, %vector_index : index - %ql_word_index0, %ql_word_index1, %qh_word_index = index.assume %ql_word_index00, %ql_word_index10, %qh_word_index0 [range(%ql_word_index00, 0, 23), range(%ql_word_index10, 8, 31), range(%qh_word_index0, 0, 15)] : index, index, index - %ql_word0 = vector.load %ql_view[%ql_word_index0] : view<32xi32> -> vector<1xi32> - %ql_word1 = vector.load %ql_view[%ql_word_index1] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> - %ql0 = vector.andi %ql_word0, %nibble_mask : vector<1xi32> - %ql1 = vector.andi %ql_word1, %nibble_mask : vector<1xi32> - %ql_word0_high = vector.shrui %ql_word0, %c4_i32v : vector<1xi32> - %ql_word1_high = vector.shrui %ql_word1, %c4_i32v : vector<1xi32> - %ql2 = vector.andi %ql_word0_high, %nibble_mask : vector<1xi32> - %ql3 = vector.andi %ql_word1_high, %nibble_mask : vector<1xi32> - %qh0_low = vector.andi %qh_word, %high0_mask : vector<1xi32> - %qh1_low = vector.andi %qh_word, %high2_mask : vector<1xi32> - %qh2 = vector.andi %qh_word, %high4_mask : vector<1xi32> - %qh3_high = vector.andi %qh_word, %high6_mask : vector<1xi32> - %qh0 = vector.shli %qh0_low, %c4_i32v : vector<1xi32> - %qh1 = vector.shli %qh1_low, %c2_i32v : vector<1xi32> - %qh3 = vector.shrui %qh3_high, %c2_i32v : vector<1xi32> - %code0 = vector.ori %ql0, %qh0 : vector<1xi32> - %code1 = vector.ori %ql1, %qh1 : vector<1xi32> - %code2 = vector.ori %ql2, %qh2 : vector<1xi32> - %code3 = vector.ori %ql3, %qh3 : vector<1xi32> - %code0_i8 = vector.bitcast %code0 : vector<1xi32> to vector<4xi8> - %code1_i8 = vector.bitcast %code1 : vector<1xi32> to vector<4xi8> - %code2_i8 = vector.bitcast %code2 : vector<1xi32> to vector<4xi8> - %code3_i8 = vector.bitcast %code3 : vector<1xi32> to vector<4xi8> - %code0_f32 = vector.uitofp %code0_i8 : vector<4xi8> to vector<4xf32> - %code1_f32 = vector.uitofp %code1_i8 : vector<4xi8> to vector<4xf32> - %code2_f32 = vector.uitofp %code2_i8 : vector<4xi8> to vector<4xf32> - %code3_f32 = vector.uitofp %code3_i8 : vector<4xi8> to vector<4xf32> - %q0 = vector.subf %code0_f32, %c32_f32v : vector<4xf32> - %q1 = vector.subf %code1_f32, %c32_f32v : vector<4xf32> - %q2 = vector.subf %code2_f32, %c32_f32v : vector<4xf32> - %q3 = vector.subf %code3_f32, %c32_f32v : vector<4xf32> - %dot0 = func.call @ggml_q6k_dot4_f32(%input0, %q0) : (vector<4xf32>, vector<4xf32>) -> (f32) - %dot1 = func.call @ggml_q6k_dot4_f32(%input1, %q1) : (vector<4xf32>, vector<4xf32>) -> (f32) - %dot2 = func.call @ggml_q6k_dot4_f32(%input2, %q2) : (vector<4xf32>, vector<4xf32>) -> (f32) - %dot3 = func.call @ggml_q6k_dot4_f32(%input3, %q3) : (vector<4xf32>, vector<4xf32>) -> (f32) - %scaled3 = scalar.mulf %dot3, %scale3 : f32 - %scaled2 = scalar.fmaf %dot2, %scale2, %scaled3 : f32 - %scaled1 = scalar.fmaf %dot1, %scale1, %scaled2 : f32 - %scaled0 = scalar.fmaf %dot0, %scale0, %scaled1 : f32 - %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %contribution0 = scalar.mulf %scaled0, %d : f32 - %contribution = scf.if %valid_block -> (f32) { - scf.yield %contribution0 : f32 - } else { - scf.yield %c0_f32 : f32 - } - func.return %contribution : f32 -} - -// Shared direct-F32 matrix-vector schedule. Target entry points below vary -// subgroup width without duplicating the Q6_K decoding or memory schedule. -func.template device @ggml_linear_q6k_f32_body(%publish_output: i1, %token_count: index, %token0: index, %pair: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : 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 - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %scale_stage_bytes = index.constant 512 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index - %row00 = index.mul %pair, %c2 : index - %row0, %launch_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index - %row1 = index.add %row0, %c1 : index - %row1_valid = index.cmp ult, %row1, %launch_output_size : index - %block_count = index.div %bounded_input_size, %c256 : index - %cohort = index.div %lane, %c16 : index - scf.if %publish_output { - %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer - %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer - %acc0, %acc1 = scf.for %block_base = [%c0 to %block_count step %c4](%row_acc0 = %c0_f32 : f32, %row_acc1 = %c0_f32 : f32) -> (f32, f32) { - %block = index.add %block_base, %cohort : index - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %row0, %block, %c2, %c0, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - %row0_input0, %row0_input1, %row0_input2, %row0_input3 = func.call @ggml_q6k_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %contribution0 = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %row0, %block, %c2, %c0, %lane, %weight_noalias, %scale_stage, %row0_input0, %row0_input1, %row0_input2, %row0_input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - %contribution1 = scf.if %row1_valid -> (f32) { - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %row1, %block, %c2, %c1, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - %row1_input0, %row1_input1, %row1_input2, %row1_input3 = func.call @ggml_q6k_load_f32_block(%launch_token_count, %bounded_input_size, %token, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %row1_contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %row1, %block, %c2, %c1, %lane, %weight_noalias, %scale_stage, %row1_input0, %row1_input1, %row1_input2, %row1_input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - scf.yield %row1_contribution : f32 - } else { - scf.yield %c0_f32 : 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 - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%launch_output_size]xf32> - scf.if %is_lane_zero { - view.store %sum0, %output_view[%token, %row0] : f32, view<[%launch_token_count]x[%launch_output_size]xf32> - } - scf.if %row1_valid { - scf.if %is_lane_zero { - view.store %sum1, %output_view[%token, %row1] : f32, view<[%launch_token_count]x[%launch_output_size]xf32> - } - } - } - func.return -} - -// Wave64 provider matching the Vulkan oracle's subgroup schedule. -kernel.def target(@ggml_q6k_gfx11_wave64) @ggml_linear_q6k_f32_wave64(%token_count: index, %input_size: index, %output_size: index) { - %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index - %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c64 = index.constant 64 : index - %padded_output_size = index.add %output_capacity, %c1 : index - %output_pairs = index.div %padded_output_size, %c2 : index - kernel.launch.config workgroups(%output_pairs, %token_capacity, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index - %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144), le(%output_size, %output_capacity)] : index - %c2 = index.constant 2 : index - %pair = kernel.workgroup.id : index - %token = kernel.workgroup.id : index - %row0 = index.mul %pair, %c2 : index - %valid_token = index.cmp ult, %token, %bounded_token_count : index - %valid_row = index.cmp ult, %row0, %bounded_output_size : index - %publish_output = scalar.andi %valid_token, %valid_row : i1 - %c0 = index.constant 0 : index - %safe_token = scf.select %valid_token, %token, %c0 : index - %safe_pair = scf.select %valid_row, %pair, %c0 : index - func.apply(%publish_output, %bounded_token_count, %safe_token, %safe_pair, %input_size, %bounded_output_size, %input, %weight, %output) : (i1, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -// Wave32 provider testing whether two subgroups hide direct-F32 decode latency -// better on targets where wave64 is not the default execution mode. -kernel.def target(@ggml_q6k_gfx11_wave32) @ggml_linear_q6k_f32_wave32(%token_count: index, %input_size: index, %output_size: index) { - %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index - %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c64 = index.constant 64 : index - %padded_output_size = index.add %output_capacity, %c1 : index - %output_pairs = index.div %padded_output_size, %c2 : index - kernel.launch.config workgroups(%output_pairs, %token_capacity, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @ggml.linear_q6k_f32.token_capacity : index - %output_capacity = config.get @ggml.linear_q6k_f32.output_capacity : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144), le(%output_size, %output_capacity)] : index - %c2 = index.constant 2 : index - %pair = kernel.workgroup.id : index - %token = kernel.workgroup.id : index - %row0 = index.mul %pair, %c2 : index - %valid_token = index.cmp ult, %token, %bounded_token_count : index - %valid_row = index.cmp ult, %row0, %bounded_output_size : index - %publish_output = scalar.andi %valid_token, %valid_row : i1 - %c0 = index.constant 0 : index - %safe_token = scf.select %valid_token, %token, %c0 : index - %safe_pair = scf.select %valid_row, %pair, %c0 : index - func.apply(%publish_output, %bounded_token_count, %safe_token, %safe_pair, %input_size, %bounded_output_size, %input, %weight, %output) : (i1, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -// Exact-representable activations make the direct-F32 and Q8_1 contraction -// paths comparable while nonzero packed bytes exercise every Q6_K field. -check.case public @ggml_linear_q6k_f32_differential_case { - %token_count = check.literal value(2) : index - %input_size = check.literal value(2048) : index - %output_size = check.literal value(9) : index - %input = check.generate.fill value(0.00390625) : tensor<2x2048xf32> - %q8_input = check.generate.fill value(0) : tensor<2x2304xi8> - %weight = check.generate.fill value(-86) : tensor<9x8x210xi8> - %expected = check.generate.fill value(0.0) : tensor<2x9xf32> - %actual_wave64 = check.generate.fill value(1.0) : tensor<2x9xf32> - %actual_wave32 = check.generate.fill value(1.0) : tensor<2x9xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x2048xf32>, tensor<2x2304xi8>) - kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %expected) : [index, index, index](index, index, index, tensor<2x2304xi8>, tensor<9x8x210xi8>, tensor<2x9xf32>) - kernel.launch @ggml_linear_q6k_f32_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %actual_wave64) : [index, index, index](index, index, index, tensor<2x2048xf32>, tensor<9x8x210xi8>, tensor<2x9xf32>) - kernel.launch @ggml_linear_q6k_f32_wave32[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %actual_wave32) : [index, index, index](index, index, index, tensor<2x2048xf32>, tensor<9x8x210xi8>, tensor<2x9xf32>) - check.expect.close actual(%actual_wave64) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<2x9xf32> - check.expect.close actual(%actual_wave32) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<2x9xf32> - check.return -} - -check.case public @ggml_linear_q6k_f32_wave64_dense_v_benchmark_case { - %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index - %input_size = check.literal value(2048) : index - %output_size = check.literal value(512) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %weight = check.generate.fill value(0) : tensor<512x8x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> - kernel.launch @ggml_linear_q6k_f32_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]x2048xf32>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> - check.return -} - -check.case public @ggml_linear_q6k_f32_wave32_dense_v_benchmark_case { - %token_count = check.param.choice values([1, 8, 32, 128, 512]) name("token_count") : index - %input_size = check.literal value(2048) : index - %output_size = check.literal value(512) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %weight = check.generate.fill value(0) : tensor<512x8x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> - kernel.launch @ggml_linear_q6k_f32_wave32[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<[%token_count]x2048xf32>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> - check.return -} - -check.benchmark<@ggml_linear_q6k_f32_differential_case> @ggml_linear_q6k_f32_differential - -check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_decode {token_count = 1} - -check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_prefill_32 {token_count = 32} - -check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_prefill_128 {token_count = 128} - -check.benchmark<@ggml_linear_q6k_f32_wave64_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave64_dense_v_prefill_512 {token_count = 512} - -check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_decode {token_count = 1} - -check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_prefill_32 {token_count = 32} - -check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_prefill_128 {token_count = 128} - -check.benchmark<@ggml_linear_q6k_f32_wave32_dense_v_benchmark_case> @ggml_linear_q6k_f32_wave32_dense_v_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom index 6f866cb48f5e..40218bb61a75 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/linear_q6k_q8_1_x4.loom @@ -13,285 +13,19 @@ // by dense and routed projections; prefill schedules can reuse the inline // packed-row primitive while staging weights across multiple activation rows. // The shared packer is linked beside this module. +template.decl @ggml.linear_q6k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) + amdgpu.target @ggml_q6k_q8_gfx1151_wave64 {subgroup_size = 64} config.decl @ggml.linear_q6k_q8_1_x4.token_capacity : %value: index where [range(%value, 1, 2048)] config.decl @ggml.linear_q6k_q8_1_x4.output_capacity : %value: index where [range(%value, 1, 262144)] -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$15: index, %input_size$16: index) launch(%token_count$17: index, %input_size$18: index, %input: buffer, %output: buffer) func.decl @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) -// Sign-extends four packed six-bit values without scalar lane extraction. -// Each byte enters with bits [5:0] populated and leaves as signed i8. -func.def inline @ggml_q6k_sign_extend_dot4(%code: vector<1xi32>) -> (vector<4xi8>) { - %c1_i32v = vector.constant 1 : vector<1xi32> - %c2_i32v = vector.constant 2 : vector<1xi32> - %low5_mask = vector.constant 522133279 : vector<1xi32> - %bit5_mask = vector.constant 538976288 : vector<1xi32> - %sign_mask = vector.constant -522133280 : vector<1xi32> - %low5 = vector.andi %code, %low5_mask : vector<1xi32> - %bit5 = vector.andi %code, %bit5_mask : vector<1xi32> - %bit6 = vector.shli %bit5, %c1_i32v : vector<1xi32> - %bit7 = vector.shli %bit5, %c2_i32v : vector<1xi32> - %high01 = vector.ori %bit5, %bit6 : vector<1xi32> - %high = vector.ori %high01, %bit7 : vector<1xi32> - %sign = vector.xori %high, %sign_mask : vector<1xi32> - %signed_i32 = vector.ori %low5, %sign : vector<1xi32> - %signed = vector.bitcast %signed_i32 : vector<1xi32> to vector<4xi8> - func.return %signed : vector<4xi8> -} - -// Decodes four adjacent values from one 32-value Q6_K group for FP16 matrix -// staging. The group and packet coordinates match the contiguous K dimension -// consumed by WMMA tiles, unlike the lane/part mapping used by the Q8_1 dot -// contraction below. -func.def inline @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 210 : offset - %qh_byte_add = index.constant 128 : offset - %scale_byte_add = index.constant 192 : offset - %d_byte_add = index.constant 208 : offset - %c4_i32v = vector.constant 4 : vector<1xi32> - %nibble_mask = vector.constant 252645135 : vector<1xi32> - %high_mask = vector.constant 50529027 : vector<1xi32> - %c32_f32v = vector.constant 32.0 : vector<4xf32> - %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %d_byte_base = index.add %block_byte_base, %d_byte_add : offset - %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %group_in_half = index.rem %bounded_group, %c4 : index - %half = index.div %bounded_group, %c4 : index - %ql_side = index.rem %group_in_half, %c2 : index - %ql_half_word_base = index.mul %half, %c16 : index - %ql_side_word_add = index.mul %ql_side, %c8 : index - %ql_word_base = index.add %ql_half_word_base, %ql_side_word_add : index - %ql_word_index = index.add %ql_word_base, %bounded_packet : index - %qh_half_word_base = index.mul %half, %c8 : index - %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index - %nibble = index.div %group_in_half, %c2 : index - %nibble_shift_index = index.mul %nibble, %c4 : index - %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 - %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> - %qh_shift_index = index.mul %group_in_half, %c2 : index - %qh_shift_i32 = index.cast %qh_shift_index : index to i32 - %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> - %scale_packet_half = index.div %bounded_packet, %c4 : index - %scale_group_base = index.mul %bounded_group, %c2 : index - %scale_index = index.add %scale_group_base, %scale_packet_half : index - %ql_word = vector.load %ql_view[%ql_word_index] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> - %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> - %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> - %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> - %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> - %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> - %code = vector.ori %ql, %qh : vector<1xi32> - %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> - %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> - %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> - %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 - %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 - %scale = scalar.sitofp %scale_i8 : i8 to f32 - %d = scalar.extf %d_f16 : f16 to f32 - %combined_scale = scalar.mulf %scale, %d : f32 - %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> - %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> - %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> - func.return %values : vector<4xf16> -} - -// Contracts both four-value packets assigned to one lane in a Q6_K block. -// QL, QH, and the block scale are shared by the two packed parts, and their -// contributions are summed before entering the row-wide recurrence. -func.def inline @ggml_q6k_q8_1_x4_block_lane(%weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q6_block: index, %lane: index) -> (f32) { - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 210 : offset - %qh_byte_add = index.constant 128 : offset - %scale_byte_add = index.constant 192 : offset - %d_byte_add = index.constant 208 : offset - %c4_i32v = vector.constant 4 : vector<1xi32> - %c0_i32v = vector.constant 0 : vector<1xi32> - %nibble_mask = vector.constant 252645135 : vector<1xi32> - %high_mask = vector.constant 808464432 : vector<1xi32> - %c0_f32 = scalar.constant 0.0 : f32 - %bounded_lane = index.assume %lane [range(%lane, 0, 31)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %d_byte_base = index.add %block_byte_base, %d_byte_add : offset - %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %lane_mod8 = index.rem %bounded_lane, %c8 : index - %lane_mod16 = index.rem %bounded_lane, %c16 : index - %lane_div16 = index.div %bounded_lane, %c16 : index - %lane_div8_in_16 = index.div %lane_mod16, %c8 : index - %lane_div4_in_16 = index.div %lane_mod16, %c4 : index - %qh_high_base = index.mul %lane_div16, %c8 : index - %qh_index0 = index.add %qh_high_base, %lane_mod8 : index - %qh_index = index.assume %qh_index0 [range(%qh_index0, 0, 15)] : index - %ql_word = vector.load %ql_view[%bounded_lane] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_index] : view<16xi32> -> vector<1xi32> - %qh_base_shift_index = index.mul %lane_div8_in_16, %c2 : index - %qh_base_shift_i32 = index.cast %qh_base_shift_index : index to i32 - %q8_block_base = index.mul %q6_block, %c8 : index - %q8_high_add = index.mul %lane_div16, %c4 : index - %q8_quadrant = index.add %q8_high_add, %lane_div8_in_16 : index - %scale_high_base = index.mul %lane_div16, %c8 : index - %scale_lane0 = index.add %scale_high_base, %lane_div4_in_16 : index - %d_f16 = view.load %d_view[0] : view<1xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %sum = scf.for %part = [%c0 to %c2 step %c1](%accumulator = %c0_f32 : f32) -> (f32) unroll { - %bounded_part = index.assume %part [range(%part, 0, 1)] : index - %part_shift_index = index.mul %bounded_part, %c4 : index - %part_shift_i32 = index.cast %part_shift_index : index to i32 - %part_shift = vector.splat %part_shift_i32 : vector<1xi32> - %ql_shifted = vector.shrui %ql_word, %part_shift : vector<1xi32> - %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> - %qh_shift_i32 = scalar.addi %qh_base_shift_i32, %part_shift_i32 : i32 - %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> - %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> - %qh_positioned = vector.shli %qh_shifted, %c4_i32v : vector<1xi32> - %qh = vector.andi %qh_positioned, %high_mask : vector<1xi32> - %code = vector.ori %ql, %qh : vector<1xi32> - %signed_weight = func.call @ggml_q6k_sign_extend_dot4(%code) : (vector<1xi32>) -> (vector<4xi8>) - %q8_part_add = index.mul %bounded_part, %c2 : index - %q8_block_part = index.add %q8_block_base, %q8_part_add : index - %q8_block = index.add %q8_block_part, %q8_quadrant : index - %q8_values, %q8_d = func.call @ggml_q8_1_x4_word(%q8_input, %q8_row_byte_base, %q8_block, %lane_mod8) : (buffer, offset, index, index) -> (vector<4xi8>, f32) - %scale_lane = index.add %scale_lane0, %part_shift_index : index - %scale_i8 = view.load %scale_view[%scale_lane] : view<16xi8> -> i8 - %scale = scalar.sitofp %scale_i8 : i8 to 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 - %scaled0 = scalar.mulf %dot_f32, %scale : f32 - %scaled1 = scalar.mulf %scaled0, %d : f32 - %contribution = scalar.mulf %scaled1, %q8_d : f32 - %next = scalar.addf %accumulator, %contribution : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Computes one lane's partial for a complete Q6_K row. Callers choose how -// rows and activations are tiled, then reduce the returned value by subgroup. -func.def inline @ggml_q6k_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, 256, 32768), mul(%input_size, 256)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %block_count = index.div %bounded_input_size, %c256 : index - %result = scf.for %block = [%c0 to %block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) { - %contribution = func.call @ggml_q6k_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 %block_acc, %contribution : f32 - scf.yield %next : f32 - } - func.return %result : f32 -} - -// Maps four Q6_K blocks across the four 16-lane partitions of a wave64. Each -// physical lane contracts the pair of virtual wave32 lanes that own the same -// packed positions in the low and high halves of one block. The subgroup -// reduction therefore combines four complete blocks per loop iteration while -// preserving the canonical wave32 unpacking and dot-product primitive. -func.def inline @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%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, 256, 32768), mul(%input_size, 256)] : index - %bounded_lane = index.assume %lane [range(%lane, 0, 63)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c16 = index.constant 16 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %block_count = index.div %bounded_input_size, %c256 : index - %padded_block_count = index.add %block_count, %c3 : index - %iteration_count = index.div %padded_block_count, %c4 : index - %block_in_iteration = index.div %bounded_lane, %c16 : index - %virtual_lane0 = index.rem %bounded_lane, %c16 : index - %virtual_lane1 = index.add %virtual_lane0, %c16 : index - %result = scf.for %iteration = [%c0 to %iteration_count step %c1](%row_acc = %c0_f32 : f32) -> (f32) { - %block_base = index.mul %iteration, %c4 : index - %block0 = index.add %block_base, %block_in_iteration : index - %valid_block = index.cmp ult, %block0, %block_count : index - %block_acc = scf.if %valid_block -> (f32) { - %low = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane0) : (buffer, offset, buffer, offset, index, index) -> (f32) - %high = func.call @ggml_q6k_q8_1_x4_block_lane(%weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %block0, %virtual_lane1) : (buffer, offset, buffer, offset, index, index) -> (f32) - %sum = scalar.addf %low, %high : f32 - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %row_acc, %block_acc : f32 - scf.yield %next : f32 - } - func.return %result : f32 -} - -// Contracts one output channel for the calling Q6_K projection kernel. The -// caller owns publication so the same canonical contraction can feed either a -// dense logits tensor or an endpoint reduction. -func.template device @ggml_linear_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer) -> (f32, index, index, i1, index, index) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c144_bytes = index.constant 144 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %q6_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset - %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset - %lane_acc = scf.if %publish_output -> (f32) { - %channel_acc = scf.if %valid_channel -> (f32) { - %value = func.call @ggml_q6k_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) - scf.yield %value : f32 - } else { - %c0_f32 = scalar.constant 0.0 : f32 - scf.yield %c0_f32 : f32 - } - scf.yield %channel_acc : f32 - } else { - %c0_f32 = scalar.constant 0.0 : f32 - scf.yield %c0_f32 : f32 - } - %dot = kernel.subgroup.reduce %lane_acc : f32 - func.return %dot, %token, %channel, %valid_channel, %launch_token_count, %bounded_output_size : f32, index, index, i1, index, index -} +func.decl @ggml_q6k_q8_1_x4_row_lane_wave64_block4(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) // Dense raw-layout baseline. A 256-thread workgroup carries eight independent // wave32 output rows so decode launches enough waves without duplicating Q6_K @@ -316,7 +50,7 @@ kernel.def export("ggml_linear_q6k_q8_1_x4") @ggml_linear_q6k_q8_1_x4(%token_cou %token0 = kernel.workgroup.id : index %valid_token = index.cmp ult, %token0, %bounded_token_count : index %safe_token = scf.select %valid_token, %token0, %c0 : index - %dot, %token, %channel, %valid_channel, %launch_token_count, %output_bound = func.apply(%valid_token, %bounded_token_count, %safe_token, %input_size, %bounded_output_size, %q8_noalias, %weight_noalias) : (i1, index, index, index, index, buffer, buffer) -> (f32, index, index, i1, index, index) + %dot, %token, %channel, %valid_channel, %launch_token_count, %output_bound = template.apply<@ggml.linear_q6k_q8_1_x4.body>(%valid_token, %bounded_token_count, %safe_token, %input_size, %bounded_output_size, %q8_noalias, %weight_noalias) : (i1, index, index, index, index, buffer, buffer) -> (f32, index, index, i1, index, index) %c0_i32 = scalar.constant 0 : i32 %c0_offset = index.constant 0 : offset %lane = kernel.subgroup.lane.id : index diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom index 854148db0a37..f0da9d9cf4dd 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/ggml/quantize_q8_1_x4.loom @@ -6,158 +6,10 @@ // f16 ds[4][2]; // per-block (scale, quantized_sum * scale) // i32 qs[4][8]; // four signed i8 values per packed word // }; -config.decl @ggml.quantize_q8_1_x4.group_capacity : %value: index where [range(%value, 1, 524288)] - -// Loads one logical 32-element Q8_1 block from the four-way physical packing. -// `row_byte_base` addresses the first x4 group for one activation row. -func.def inline @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) { - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %group = index.div %q8_block, %c4 : index - %inner0 = index.rem %q8_block, %c4 : index - %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index - %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset - %group_byte_base = index.add %row_byte_base, %group_byte_add : offset - %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset - %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> - %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> - %d_index = index.mul %inner, %c2 : index - %s_index = index.add %d_index, %c1 : index - %word_index = index.mul %inner, %c8 : index - %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 - %s_f16 = view.load %ds_view[%s_index] : view<8xf16> -> f16 - %words = vector.load %payload_view[%word_index] : view<32xi32> -> vector<8xi32> - %values = vector.bitcast %words : vector<8xi32> to vector<32xi8> - %d = scalar.extf %d_f16 : f16 to f32 - %s = scalar.extf %s_f16 : f16 to f32 - func.return %values, %d, %s : vector<32xi8>, f32, f32 -} +template.decl @ggml.quantize_q8_1_x4.group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) -// Loads one packed four-value word and its block scale. Dot-product kernels -// use this narrower form so a lane does not load the other seven words owned -// by its subgroup peers. -func.def inline @ggml_q8_1_x4_word(%q8_input: buffer, %row_byte_base: offset, %q8_block: index, %word_in_block: index) -> (vector<4xi8>, f32) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %group = index.div %q8_block, %c4 : index - %inner0 = index.rem %q8_block, %c4 : index - %inner = index.assume %inner0 [range(%inner0, 0, 3)] : index - %word0 = index.assume %word_in_block [range(%word_in_block, 0, 7)] : index - %group_byte_add = index.scale %group, %group_bytes : index, offset -> offset - %group_byte_base = index.add %row_byte_base, %group_byte_add : offset - %payload_byte_base = index.add %group_byte_base, %payload_byte_add : offset - %ds_view = buffer.view %q8_input[%group_byte_base] : buffer -> view<8xf16> - %payload_view = buffer.view %q8_input[%payload_byte_base] : buffer -> view<32xi32> - %d_index = index.mul %inner, %c2 : index - %inner_word_base = index.mul %inner, %c8 : index - %word_index = index.add %inner_word_base, %word0 : index - %d_f16 = view.load %ds_view[%d_index] : view<8xf16> -> f16 - %packed = vector.load %payload_view[%word_index] : view<32xi32> -> vector<1xi32> - %values = vector.bitcast %packed : vector<1xi32> to vector<4xi8> - %d = scalar.extf %d_f16 : f16 to f32 - func.return %values, %d : vector<4xi8>, f32 -} +config.decl @ggml.quantize_q8_1_x4.group_capacity : %value: index where [range(%value, 1, 524288)] -// Packs one explicit physical group. Callers provide the complete physical -// group domain and an in-range ordinal, and guarantee that one complete -// 32-lane wave executes after all 128 source values are visible. The uniform -// publication predicate suppresses only destination writes so every lane still -// reaches the workgroup barriers. -func.template device @ggml_quantize_q8_1_x4_group_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { - %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index - %lane = kernel.workitem.id : 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 - %c32 = index.constant 32 : index - %c128 = index.constant 128 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 128 : offset - %scratch_bytes = index.constant 144 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_offset = index.constant 0 : offset - %launched_element_count = index.mul %group_count, %c128 : index - %block_in_group0 = index.div %lane, %c8 : index - %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index - %word_in_block0 = index.rem %lane, %c8 : index - %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index - %group_element_base = index.mul %group, %c128 : index - %block_element_add = index.mul %block_in_group, %c32 : index - %block_word_add = index.mul %block_in_group, %c8 : index - %word_element_add = index.mul %word_in_block, %c4 : index - %input_block_base = index.add %group_element_base, %block_element_add : index - %input_index = index.add %input_block_base, %word_element_add : index - %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> - %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> - %absolute_values = vector.absf %input_values : vector<4xf32> - %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<32xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<4xf32> - view.store %thread_max, %scratch_values[%lane] : f32, view<32xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index - scf.if %is_cohort_leader { - %cohort_base = index.mul %block_in_group, %c8 : index - %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> - %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%block_in_group] : f32, view<4xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %d = view.load %scratch_d[%block_in_group] : view<4xf32> -> f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> - %packed_word_index0 = index.add %block_word_add, %word_in_block : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - scf.if %publish_output { - vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> - } - %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_sum, %scratch_values[%lane] : f32, view<32xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %publishes_metadata = scalar.andi %is_cohort_leader, %publish_output : i1 - scf.if %publishes_metadata { - %cohort_base = index.mul %block_in_group, %c8 : index - %cohort_sums = vector.load %scratch_values[%cohort_base] : view<32xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - %s_index = index.add %ds_index, %c1 : index - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - func.return -} // A 32-lane workgroup owns one physical group. Each eight-lane cohort loads // four values, computes one Q8_1 block, and writes disjoint metadata and packed @@ -180,7 +32,7 @@ kernel.def @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) { %valid_group = index.cmp ult, %group0, %group_count : index %safe_group = scf.select %valid_group, %group0, %c0 : index %group = index.assume %safe_group [lt(%safe_group, %group_count)] : index - func.apply(%valid_group, %group_count, %group, %input, %output) : (i1, index, index, buffer, buffer) + template.apply<@ggml.quantize_q8_1_x4.group_body>(%valid_group, %group_count, %group, %input, %output) : (i1, index, index, buffer, buffer) 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 643fdf19fe95..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 @@ -1,411 +1,124 @@ { - "schema": "ggml-hrx-qwen-kernel-corpus-v2", - "upstream_repository": "https://github.com/rocm/hrx", - "upstream_revision": "d343e4ad5167dca6758cb0d2db08785e43e99070", - "source_subdirectory": "experimental/qwen_moe/kernels", - "qwen_endpoint_source_subdirectory": "experimental/qwen/kernels", - "corpus_sha256": "ce2711ba360db150344d63658b6fc0024954e4ad4c598233caed7680ca85f993", - "upstream_corpus_sha256": "76f504d2c19e85a1b51679bb2edc5e9a60b865f545ecaabacedd14ba7db15842", - "owned_corpus_sha256": "1126fc55f6123c13ee9a8a8d146082bd9b63d9fd1ed123e27478ff964e58ad44", - "build_bazel_sha256": "598eeb34f8606182a732299c989387dd450c0d8f33e4275635b4a679c0f986d0", + "schema": "ggml-hrx-kernel-corpus-v1", + "upstream_revision": "local", "files": [ { - "path": "ggml/linear_q6k_f32.loom", - "sha256": "f20c9cd0ce58c6158b8b45bf6c9502dc22d20d2a67eee975d96e7625f87ef615", - "size": 28513 + "path": "../loom-libs/motifs/q6_k_f16.loom" }, { - "path": "ggml/linear_q6k_q8_1_x4.loom", - "sha256": "865795fdd3111e1e29bfe07325a3057c55656773109caec1224ac42e03d86bba", - "size": 32592 + "path": "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom" }, { - "path": "ggml/quantize_q8_1_x4.loom", - "sha256": "a9e6e995dd541142574058d4412f313beb87b6c6f4bb750b23a0a3047c951183", - "size": 16711 + "path": "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom" }, { - "path": "qwen3_moe/attention_postprocess_f32_f16.loom", - "sha256": "9125cd1ef87cb0d0ca8f4f6fc0c0a5bec29f9d6b30894287b53ad2d98c9433fb", - "size": 25955 + "path": "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom" }, { - "path": "qwen3_moe/attention_prepare_quantized.loom", - "sha256": "4bd47e89d1569e3622de60d88a6e09309afe35d48d3b6b1f25469e3b71aa7118", - "size": 29030 + "path": "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom" }, { - "path": "qwen3_moe/attention_qkv_postprocess_fused.loom", - "sha256": "409f9f95ad8b0a58691cc74d665489d3d65ae37c73fb682b405ebff0bf1c0c79", - "size": 24797 + "path": "../loom-libs/motifs/routing/router_projection_storage.loom" }, { - "path": "qwen3_moe/attention_qkv_quantized.loom", - "sha256": "7b631b54e042ead32c52b8ddacba4fd58de588571b586a05719f0ddcee080cc1", - "size": 25990 + "path": "../loom-libs/motifs/routing/router_topk_row.loom" }, { - "path": "qwen3_moe/attention_qkv_same_format_prefill.loom", - "sha256": "55aeb3857c0646e9a1f05570ae6860fe42b0cfa91198868fc46bbed19f1c02fb", - "size": 12730 + "path": "../loom-libs/motifs/routing/router_projection_fused.loom" }, { - "path": "qwen3_moe/batched_decode_expert_dispatch.loom", - "sha256": "d97e2590dff046b79248b4a4b64828a5d2a328488ef42d4db7f489f0b2fdc8e6", - "size": 43960 + "path": "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" }, { - "path": "qwen3_moe/batched_decode_gate_up_q4k.loom", - "sha256": "292074ba00372ee7fe409047e54c7d98b6a5bb2436e2245d993f50d124a55307", - "size": 27198 + "path": "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom" }, { - "path": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "sha256": "331fbf77a24fd22424d75e7d7716cc78af8f43899bfe15b2ee947aa9c68f0bd8", - "size": 68289 + "path": "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" }, { - "path": "qwen3_moe/expert_table_partition_fused.loom", - "sha256": "829eb55d5aff318c6872f0270fbad767748965a9943a229f6b17657b2a9b972d", - "size": 17122 + "path": "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom" }, { - "path": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", - "sha256": "e02cb287df33b22d103117ce099310d9650061215b7df86f10b3ec3df5bf5e42", - "size": 55705 + "path": "../loom-libs/motifs/quantize_q8_1_x4.loom" }, { - "path": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", - "sha256": "4a44c7d01dd37802656e250cae410593cc4cfb87e2d05ff6152bb10847e29903", - "size": 19360 + "path": "../loom-libs/motifs/quantize_q8_1_x4_group.loom" }, { - "path": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "sha256": "466b17c488f5c0889aa55c208894a67eb6db0e1d91ba0e08a5b3e8fb4c82a94f", - "size": 94449 + "path": "ggml/linear_q6k_q8_1_x4.loom" }, { - "path": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", - "sha256": "8bf53905b313dbe8ee6007e7dc7deac64d3b7412bf36180cb19c3c66d8a3255b", - "size": 9162 + "path": "ggml/quantize_q8_1_x4.loom" }, { - "path": "qwen3_moe/flash_attention_f32_f16_wmma.loom", - "sha256": "f5a745cee4cdc8d8287ce42398050a501f9f8b0c984654c4e53a8e89eba8c4c8", - "size": 71377 + "path": "qwen3_moe/attention_postprocess_f32_f16.loom" }, { - "path": "qwen3_moe/model_config.loom", - "sha256": "802d51e2d035461beb185cff93be4d602f98b408e385e0284b1096fe259a6095", - "size": 960 + "path": "qwen3_moe/attention_prepare_quantized.loom" }, { - "path": "qwen3_moe/routed_down_q4k.loom", - "sha256": "cf7027093e3baadd2e6cc50ca70943c7f228781d0e042c7673110ec4e5dd9e29", - "size": 33238 + "path": "qwen3_moe/attention_qkv_postprocess_fused.loom" }, { - "path": "qwen3_moe/routed_down_q6k.loom", - "sha256": "fd8bc76bafd174fd57e298849f65c2876186b1c6aeaa90e3b419f26b2a5d0326", - "size": 64580 + "path": "qwen3_moe/attention_qkv_quantized.loom" }, { - "path": "qwen3_moe/routed_down_next_q8.loom", - "sha256": "3eccaa3290b83967a44958064bbc11960f9db77e68436178c2c5a50052014e2b", - "size": 4239 + "path": "qwen3_moe/dense_linear_quantized_f16_wmma.loom" }, { - "path": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "sha256": "b97a8e15a581acbc9b1de5776c7ef2ac727e8c138686ed496e2913814d2fd343", - "size": 75125 + "path": "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom" }, { - "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", - "sha256": "351f7bcf6230058df9b78f59abb7011cf5690eedf5862ec0b75b537bf27e0e52", - "size": 11933 + "path": "qwen3_moe/motifs/attention_prepare_target.loom" }, { - "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom", - "sha256": "1ce06fd438cdebd63805b06066445cee717800ad1228ccfdc8294e8faccd6d95", - "size": 19387 + "path": "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom" }, { - "path": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "sha256": "e5d5ed276070ae25e8f1d915e927716da519314d34319db4b722d30589233deb", - "size": 125922 + "path": "qwen3_moe/expert_table_partition_fused.loom" }, { - "path": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "sha256": "77f024758c30fd5c834f1429049021af3616a7162486115a41735aa81bf4d613", - "size": 73406 + "path": "qwen3_moe/motifs/expert_table_partition_prefill_512.loom" }, { - "path": "qwen3_moe/router_projection_f32.loom", - "sha256": "6c3eaf6834af0a58ecbd91702e009d410c4e7d52eac106a98a0f634f1d452ed5", - "size": 18655 + "path": "qwen3_moe/model_config.loom" }, { - "path": "qwen3_moe/router_projection_top8_fused_f32.loom", - "sha256": "193615e54d175a8d2231bfda1f2c7672e338c1e0c4eef01cf0a5c3e35b26daa4", - "size": 17578 + "path": "qwen3_moe/routed_down_quantized_f16_wmma.loom" }, { - "path": "qwen3_moe/router_top8_f32.loom", - "sha256": "84738d2baee5b1a54dff2939a933024222fe8dee76ea30d71af1921b23a9e665", - "size": 14431 + "path": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" }, { - "path": "../qwen_owned/token_embedding_q4k.loom", - "sha256": "5cb5eed0b07991e1af246965bfeff0034b066d1645d5d38cbcdd63757dee5c9a", - "size": 12428, - "upstream_path": "experimental/qwen/kernels/token_embedding_q4k.loom" + "path": "../loom-libs/motifs/matmul_id/vector_pair.loom" }, { - "path": "../qwen_owned/attention_metadata.loom", - "sha256": "9614108d1680e6d0d86187c75ce0505e44fb1f53f3117f38baf56ab686b0632c", - "size": 13620, - "upstream_path": "experimental/qwen/kernels/attention_metadata.loom" - }, - { - "path": "../qwen_owned/token_embedding_bringup_workaround.loom", - "sha256": "13666225ffea241039106b33b6ed26000f9cc271b50249f5bdce61a44243e417", - "size": 15073, - "owner": "ggml-hrx" + "path": "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom" }, { - "path": "../qwen_owned/attention_state_initialize.loom", - "sha256": "e0492f884ea2bc0f0818e77d43b6639a81fdf994a1940416abaeac86e4165f55", - "size": 7222, - "owner": "ggml-hrx" + "path": "../loom-libs/motifs/elementwise/binary_apply.loom" }, { - "path": "../qwen_owned/attention_metadata_bringup_workaround.loom", - "sha256": "44a462abbda7e68dc84becc45501885d834b20d8445be1eb0bc7690d829b0935", - "size": 12346, - "owner": "ggml-hrx" + "path": "../loom-libs/ops/matmul_id/vector_pair_binary.loom" }, { - "path": "../hrx_owned/gather_add_f32.loom", - "sha256": "c0f28e67b0f26d49664c1a307afb66367f3f05663e939953fba36885b7a2f6ac", - "size": 4497, - "owner": "ggml-hrx" + "path": "../loom-libs/ops/moe_routing_tables.loom" }, { - "path": "../hrx_owned/add_f32.loom", - "sha256": "249b1d4c7167d0533277481cbb0e930207f14f9e2ba75fc1b8754eb0717ed094", - "size": 2444, - "owner": "ggml-hrx" - } - ], - "exports": [ - { - "name": "ggml_add_f32", - "symbol": "ggml_add_f32", - "source": "../hrx_owned/add_f32.loom", - "workload_parameters": [ - { - "name": "element_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "element_count", - "type": "index" - } - ], - "bindings": [ - "a", - "b", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../hrx_owned/add_f32.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] + "path": "qwen3_moe/router_projection_f32.loom" }, { - "name": "ggml_gather_add_f32", - "symbol": "ggml_gather_add_f32", - "source": "../hrx_owned/gather_add_f32.loom", - "workload_parameters": [ - { - "name": "source_token_count", - "type": "index" - }, - { - "name": "output_token_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "source_token_count", - "type": "index" - }, - { - "name": "output_token_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ], - "bindings": [ - "attention", - "residual", - "output_ids", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../hrx_owned/gather_add_f32.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] + "path": "qwen3_moe/router_projection_top8_fused_f32.loom" }, { - "name": "ggml_linear_q6k_f32_wave32", - "symbol": "ggml_linear_q6k_f32_wave32", - "source": "ggml/linear_q6k_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "linear_q6k_f32_linked", - "primary_sources": [ - "ggml/linear_q6k_f32.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" - ] + "path": "qwen3_moe/router_top8_f32.loom" }, { - "name": "ggml_linear_q6k_f32_wave64", - "symbol": "ggml_linear_q6k_f32_wave64", - "source": "ggml/linear_q6k_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "linear_q6k_f32_linked", - "primary_sources": [ - "ggml/linear_q6k_f32.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, + "path": "../qwen/attention_metadata.loom", + "upstream_path": "experimental/qwen/kernels/attention_metadata.loom" + } + ], + "exports": [ { "name": "ggml_linear_q6k_q8_1_x4", "symbol": "ggml_linear_q6k_q8_1_x4", @@ -456,11 +169,17 @@ "ggml/linear_q6k_q8_1_x4.loom" ], "library_sources": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -513,41 +232,19 @@ "ggml/linear_q6k_q8_1_x4.loom" ], "library_sources": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, - { - "name": "ggml_q8_1_x4_inspect_one_group", - "symbol": "ggml_q8_1_x4_inspect_one_group", - "source": "ggml/quantize_q8_1_x4.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "packed", - "words", - "d_values", - "s_values" - ], - "binding_access": [ - "read", - "write", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "ggml/quantize_q8_1_x4.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, { "name": "ggml_quantize_q8_1_x4_f32", "symbol": "ggml_quantize_q8_1_x4_f32", @@ -586,61 +283,12 @@ "primary_sources": [ "ggml/quantize_q8_1_x4.loom" ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_attention_key_q4_aggregate_prefill_512", - "symbol": "qwen3_moe_attention_key_q4_aggregate_prefill_512", - "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "combined_weight", - "combined_output" - ], - "binding_access": [ - "read", - "read", - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "attention_qkv_same_format_prefill_linked", - "primary_sources": [ - "qwen3_moe/attention_qkv_same_format_prefill.loom" - ], "library_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { @@ -703,11 +351,13 @@ "qwen3_moe/attention_postprocess_f32_f16.loom" ], "library_sources": [ - "qwen3_moe/model_config.loom" + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ] }, "compile_dependencies": [ - "qwen3_moe/model_config.loom" + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ] }, { @@ -780,180 +430,124 @@ "qwen3_moe/attention_qkv_postprocess_fused.loom" ], "library_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_attention_qkv_postprocess_fused_decode_q4", - "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode_q4", - "source": "qwen3_moe/attention_qkv_postprocess_fused.loom", + "name": "qwen3_moe_build_expert_table_partition_prefill_512", + "symbol": "qwen3_moe_build_expert_table_partition_prefill_512", + "source": "qwen3_moe/motifs/expert_table_partition_prefill_512.loom", "workload_parameters": [ { "name": "token_count", "type": "index" }, { - "name": "cache_row_count", + "name": "route_count", "type": "index" - } - ], - "launch_parameters": [ + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", + "type": "index" + } + ], + "launch_parameters": [ { "name": "token_count", "type": "index" }, { - "name": "cache_row_count", + "name": "route_count", + "type": "index" + }, + { + "name": "route_stride", + "type": "index" + }, + { + "name": "expert_count", "type": "index" } ], "bindings": [ - "q8_input", - "query_weight", - "key_weight", - "value_weight", - "positions", - "key_cache_indices", - "value_cache_indices", - "query_output_raw", - "key_output_raw", - "value_output_raw", - "query_norm_weight", - "key_norm_weight", - "inverse_frequencies", - "query_output", - "key_cache", - "value_cache", - "completion_counters" + "route_ids", + "expert_table", + "partition_table", + "completion_counter" ], "binding_access": [ "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read_write", - "read_write", - "read_write", + "write", + "write", "read_write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_qkv_postprocess_fused_linked", + "link_module": "expert_table_partition_fused_linked", "primary_sources": [ - "qwen3_moe/attention_qkv_postprocess_fused.loom" + "qwen3_moe/motifs/expert_table_partition_prefill_512.loom" ], "library_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_attention_qkv_postprocess_fused_decode_q6", - "symbol": "qwen3_moe_attention_qkv_postprocess_fused_decode_q6", - "source": "qwen3_moe/attention_qkv_postprocess_fused.loom", + "name": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", + "source": "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom", "workload_parameters": [ { "name": "token_count", "type": "index" - }, - { - "name": "cache_row_count", - "type": "index" } ], "launch_parameters": [ { "name": "token_count", "type": "index" - }, - { - "name": "cache_row_count", - "type": "index" } ], "bindings": [ "q8_input", - "query_weight", - "key_weight", - "value_weight", - "positions", - "key_cache_indices", - "value_cache_indices", - "query_output_raw", - "key_output_raw", - "value_output_raw", - "query_norm_weight", - "key_norm_weight", - "inverse_frequencies", - "query_output", - "key_cache", - "value_cache", - "completion_counters" + "weight", + "output", + "norm_weight", + "normalized_output", + "completion_counter", + "next_q8_output" ], "binding_access": [ - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", - "read", "read", "read", "read_write", + "read", "read_write", "read_write", "read_write" @@ -961,38 +555,30 @@ "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_qkv_postprocess_fused_linked", + "link_module": "dense_linear_quantized_f16_wmma_linked", "primary_sources": [ - "qwen3_moe/attention_qkv_postprocess_fused.loom" + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom" ], "library_sources": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_attention_qkv_q4_prefill_512", - "symbol": "qwen3_moe_attention_qkv_q4_prefill_512", - "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "name": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", + "source": "qwen3_moe/attention_prepare_quantized.loom", "workload_parameters": [ { "name": "token_count", @@ -1007,45 +593,43 @@ ], "bindings": [ "input", - "combined_weight", - "combined_output" + "weight", + "normalized_output", + "q8_output" ], "binding_access": [ "read", "read", - "read" + "read_write", + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_qkv_same_format_prefill_linked", + "link_module": "attention_prepare_quantized_linked", "primary_sources": [ - "qwen3_moe/attention_qkv_same_format_prefill.loom" + "qwen3_moe/attention_prepare_quantized.loom" ], "library_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, "compile_dependencies": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "qwen3_moe_attention_qkv_quantized", - "symbol": "qwen3_moe_attention_qkv_quantized", - "source": "qwen3_moe/attention_qkv_quantized.loom", + "name": "qwen3_moe_routed_down_weighted_reduce_f16_f32", + "symbol": "qwen3_moe_routed_down_weighted_reduce_f16_f32", + "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", "workload_parameters": [ { "name": "token_count", @@ -1059,54 +643,36 @@ } ], "bindings": [ - "q8_input", - "query_weight", - "key_weight", - "value_weight", - "query_output", - "key_output", - "value_output" + "route_weights", + "routed_output", + "residual_input", + "output" ], "binding_access": [ "read", "read", "read", - "read", - "read_write", - "read_write", - "read_write" + "write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_qkv_quantized_linked", + "link_module": "routed_down_quantized_f16_wmma_linked", "primary_sources": [ - "qwen3_moe/attention_qkv_quantized.loom" + "qwen3_moe/routed_down_quantized_f16_wmma.loom" ], "library_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" ] }, "compile_dependencies": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" ] }, { - "name": "qwen3_moe_attention_query_q4_aggregate_prefill_512", - "symbol": "qwen3_moe_attention_query_q4_aggregate_prefill_512", - "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", + "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", + "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -1120,46 +686,44 @@ } ], "bindings": [ - "input", - "combined_weight", - "combined_output" + "route_weights", + "routed_output", + "hidden_state", + "next_norm_weight", + "next_projection_input" ], "binding_access": [ "read", "read", - "read" + "read_write", + "read", + "write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_qkv_same_format_prefill_linked", + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", "primary_sources": [ - "qwen3_moe/attention_qkv_same_format_prefill.loom" + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "library_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.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": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.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" ] }, { - "name": "qwen3_moe_attention_rmsnorm_quantize_q8_1_x4", - "symbol": "qwen3_moe_attention_rmsnorm_quantize_q8_1_x4", - "source": "qwen3_moe/attention_prepare_quantized.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", @@ -1173,36 +737,46 @@ } ], "bindings": [ - "input", - "weight", - "output" + "route_weights", + "routed_output", + "hidden_state", + "next_norm_weight", + "next_projection_input", + "q8_output" ], "binding_access": [ "read", "read", - "read_write" + "read_write", + "read", + "write", + "write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_prepare_quantized_linked", + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.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": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.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" ] }, { - "name": "qwen3_moe_attention_value_q4_aggregate_prefill_512", - "symbol": "qwen3_moe_attention_value_q4_aggregate_prefill_512", - "source": "qwen3_moe/attention_qkv_same_format_prefill.loom", + "name": "qwen3_moe_router_projection_f32_four_row_wave32", + "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", + "source": "qwen3_moe/router_projection_f32.loom", "workload_parameters": [ { "name": "token_count", @@ -1217,218 +791,103 @@ ], "bindings": [ "input", - "combined_weight", - "combined_output" + "weight", + "output" ], "binding_access": [ "read", "read", - "read" + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "attention_qkv_same_format_prefill_linked", + "link_module": "router_projection_f32_linked", "primary_sources": [ - "qwen3_moe/attention_qkv_same_format_prefill.loom" + "qwen3_moe/router_projection_f32.loom" ], "library_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/routing/router_projection_storage.loom", + "qwen3_moe/model_config.loom" ] }, "compile_dependencies": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/routing/router_projection_storage.loom", + "qwen3_moe/model_config.loom" ] }, { - "name": "qwen3_moe_batched_decode_gate_up_q4k_rows2", - "symbol": "qwen3_moe_batched_decode_gate_up_q4k_rows2", - "source": "qwen3_moe/batched_decode_gate_up_q4k.loom", + "name": "qwen3_moe_router_projection_top8_fused_decode_f32", + "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", + "source": "qwen3_moe/router_projection_top8_fused_f32.loom", "workload_parameters": [ - { - "name": "descriptor_count", - "type": "index" - }, - { - "name": "queue_ordinal", - "type": "index" - }, { "name": "token_count", "type": "index" }, { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", + "name": "route_id_stride", "type": "index" } ], "launch_parameters": [ - { - "name": "descriptor_count", - "type": "index" - }, - { - "name": "queue_ordinal", - "type": "index" - }, { "name": "token_count", "type": "index" }, { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "queue_descriptors", - "assignment_ordinals", - "q8_input", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "batched_decode_gate_up_q4k_linked", - "primary_sources": [ - "qwen3_moe/batched_decode_gate_up_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/batched_decode_expert_dispatch.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/batched_decode_expert_dispatch.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ] - }, - { - "name": "qwen3_moe_build_batched_decode_expert_dispatch", - "symbol": "qwen3_moe_build_batched_decode_expert_dispatch", - "source": "qwen3_moe/batched_decode_expert_dispatch.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", + "name": "route_id_stride", "type": "index" } ], "bindings": [ + "input", + "weight", + "logits", + "completion_counter", "route_ids", - "assignment_ordinals", - "queue_counts", - "queue_descriptors" + "route_weights" ], "binding_access": [ "read", "read", - "read", - "read" + "read_write", + "read_write", + "read_write", + "read_write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "batched_decode_expert_dispatch_linked", + "link_module": "router_projection_top8_fused_f32_linked", "primary_sources": [ - "qwen3_moe/batched_decode_expert_dispatch.loom" + "qwen3_moe/router_projection_top8_fused_f32.loom" ], "library_sources": [ - "qwen3_moe/model_config.loom" + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ] }, "compile_dependencies": [ - "qwen3_moe/model_config.loom" + "qwen3_moe/model_config.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ] }, { - "name": "qwen3_moe_build_batched_decode_expert_dispatch_reference", - "symbol": "qwen3_moe_build_batched_decode_expert_dispatch_reference", - "source": "qwen3_moe/batched_decode_expert_dispatch.loom", + "name": "qwen3_moe_router_top8_f32", + "symbol": "qwen3_moe_router_top8_f32", + "source": "qwen3_moe/router_top8_f32.loom", "workload_parameters": [ { "name": "token_count", "type": "index" }, { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", + "name": "route_id_stride", "type": "index" } ], @@ -1438,119 +897,48 @@ "type": "index" }, { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", + "name": "route_id_stride", "type": "index" } ], "bindings": [ + "logits", "route_ids", - "assignment_ordinals", - "queue_counts", - "queue_descriptors" + "route_weights" ], "binding_access": [ "read", - "read", - "read", - "read" + "write", + "write" ], "target_selector": "", "compile_recipe": { "mode": "archive", - "link_module": "batched_decode_expert_dispatch_linked", + "link_module": "router_top8_f32_linked", "primary_sources": [ - "qwen3_moe/batched_decode_expert_dispatch.loom" + "qwen3_moe/router_top8_f32.loom" ], "library_sources": [ + "../loom-libs/motifs/routing/router_topk_row.loom", "qwen3_moe/model_config.loom" ] }, "compile_dependencies": [ + "../loom-libs/motifs/routing/router_topk_row.loom", "qwen3_moe/model_config.loom" ] }, { - "name": "qwen3_moe_build_expert_partition_table", - "symbol": "qwen3_moe_build_expert_partition_table", - "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ], - "bindings": [ - "expert_table", - "partition_table" - ], - "binding_access": [ - "read", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_build_expert_table", - "symbol": "qwen3_moe_build_expert_table", - "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", + "name": "qwen_attention_metadata", + "symbol": "qwen_attention_metadata", + "source": "../qwen/attention_metadata.loom", "workload_parameters": [ { "name": "token_count", "type": "index" }, { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", + "name": "context_capacity", "type": "index" } ], @@ -1560,3165 +948,186 @@ "type": "index" }, { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", + "name": "context_capacity", "type": "index" } ], "bindings": [ - "route_ids", - "expert_table" + "control", + "positions", + "key_cache_indices", + "value_cache_indices", + "attention_mask" ], "binding_access": [ "read", - "write" + "read_write", + "read_write", + "read_write", + "read_write" ], "target_selector": "", "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_linked", + "mode": "direct", "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_build_expert_table_partition_prefill_512", - "symbol": "qwen3_moe_build_expert_table_partition_prefill_512", - "source": "qwen3_moe/expert_table_partition_fused.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - } - ], - "bindings": [ - "route_ids", - "expert_table", - "partition_table", - "completion_counter" - ], - "binding_access": [ - "read", - "write", - "write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "expert_table_partition_fused_linked", - "primary_sources": [ - "qwen3_moe/expert_table_partition_fused.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_dense_linear_q4k_f16_wmma", - "symbol": "qwen3_moe_dense_linear_q4k_f16_wmma", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_dense_linear_q4k_f16_wmma_parameterized", - "symbol": "qwen3_moe_dense_linear_q4k_f16_wmma_parameterized", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - }, - { - "name": "output_accumulation", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - }, - { - "name": "output_accumulation", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_dense_linear_q4k_q8_1_x4", - "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", - "symbol": "qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "weight", - "output", - "norm_weight", - "normalized_output", - "completion_counter", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_dense_linear_q6k_f16_wmma", - "symbol": "qwen3_moe_dense_linear_q6k_f16_wmma", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_dense_linear_q6k_f16_wmma_parameterized", - "symbol": "qwen3_moe_dense_linear_q6k_f16_wmma_parameterized", - "source": "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - }, - { - "name": "output_accumulation", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - }, - { - "name": "output_accumulation", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "dense_linear_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" - ], - "library_sources": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_flash_attention_decode_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_decode_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "query", - "key", - "value", - "mask", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "query", - "key", - "value", - "mask", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "query", - "key", - "value", - "mask", - "partial_max", - "partial_sum", - "partial_output", - "completion_counter", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read_write", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8", - "symbol": "qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8", - "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "query", - "key", - "value", - "mask", - "partial_max", - "partial_sum", - "partial_output", - "completion_counter", - "output", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read_write", - "read_write", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_mask_513_of_576", - "symbol": "qwen3_moe_flash_attention_decode_split_mask_513_of_576", - "source": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "mask" - ], - "binding_access": [ - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" - ], - "library_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_pack_completed_q8_test", - "symbol": "qwen3_moe_flash_attention_decode_split_pack_completed_q8_test", - "source": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "input", - "output" - ], - "binding_access": [ - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" - ], - "library_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "query", - "key", - "value", - "mask", - "partial_max", - "partial_sum", - "partial_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_quantize_reference_4096", - "symbol": "qwen3_moe_flash_attention_decode_split_quantize_reference_4096", - "source": "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "input", - "output" - ], - "binding_access": [ - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" - ], - "library_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_flash_attention_decode_split_reduce_f32", - "symbol": "qwen3_moe_flash_attention_decode_split_reduce_f32", - "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "partial_max", - "partial_sum", - "partial_output", - "output" - ], - "binding_access": [ - "read_write", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_f32_f16_wmma", - "symbol": "qwen3_moe_flash_attention_f32_f16_wmma", - "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "key_value_token_count", - "type": "index" - } - ], - "bindings": [ - "query", - "key", - "value", - "mask", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_test_extract_row", - "symbol": "qwen3_moe_flash_attention_test_extract_row", - "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom", - "workload_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "context_count", - "type": "index" - }, - { - "name": "source_row", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "query_token_count", - "type": "index" - }, - { - "name": "context_count", - "type": "index" - }, - { - "name": "source_row", - "type": "index" - } - ], - "bindings": [ - "source_query", - "source_mask", - "target_query", - "target_mask" - ], - "binding_access": [ - "read", - "read", - "read", - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_flash_attention_test_make_causal_mask", - "symbol": "qwen3_moe_flash_attention_test_make_causal_mask", - "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "mask" - ], - "binding_access": [ - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "qwen3_moe/flash_attention_f32_f16_wmma.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen3_moe_rmsnorm_f32", - "symbol": "qwen3_moe_rmsnorm_f32", - "source": "qwen3_moe/attention_prepare_quantized.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": "archive", - "link_module": "attention_prepare_quantized_linked", - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "symbol": "qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", - "source": "qwen3_moe/attention_prepare_quantized.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "normalized_output", - "q8_output" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "attention_prepare_quantized_linked", - "primary_sources": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check", - "symbol": "qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "input", - "norm_weight", - "q8_output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q4k_f16_wmma_grouped", - "symbol": "qwen3_moe_routed_down_q4k_f16_wmma_grouped", - "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q4k_q8_1_x4", - "symbol": "qwen3_moe_routed_down_q4k_q8_1_x4", - "source": "qwen3_moe/routed_down_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q4k.loom" - ], - "library_sources": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q4k_q8_1_x4_next_q8", - "symbol": "qwen3_moe_routed_down_q4k_q8_1_x4_next_q8", - "source": "qwen3_moe/routed_down_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output", - "norm_weight", - "completion_counter", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q4k.loom" - ], - "library_sources": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_f16_wmma_grouped", - "symbol": "qwen3_moe_routed_down_q6k_f16_wmma_grouped", - "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_f32_wave64", - "symbol": "qwen3_moe_routed_down_q6k_f32_wave64", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "route_ids", - "route_weights", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "symbol": "qwen3_moe_routed_down_q6k_f32_wave64_next_q8", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "input", - "route_ids", - "route_weights", - "weight", - "output", - "norm_weight", - "completion_counter", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_q8_1_x4", - "symbol": "qwen3_moe_routed_down_q6k_q8_1_x4", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_q6k_q8_1_x4_next_q8", - "symbol": "qwen3_moe_routed_down_q6k_q8_1_x4_next_q8", - "source": "qwen3_moe/routed_down_q6k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "input_size", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "route_weights", - "weight", - "output", - "norm_weight", - "completion_counter", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_q6k_linked", - "primary_sources": [ - "qwen3_moe/routed_down_q6k.loom" - ], - "library_sources": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - "compile_dependencies": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_weighted_reduce_f16_f32", - "symbol": "qwen3_moe_routed_down_weighted_reduce_f16_f32", - "source": "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "route_weights", - "routed_output", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_quantized_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32", - "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" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read", - "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": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4", - "symbol": "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4", - "source": "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.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" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked", - "primary_sources": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" - ], - "library_sources": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "partition_table", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8", - "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8", - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8", - "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "route_stride", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "route_ids", - "gate_weight", - "up_weight", - "output", - "completion_counters", - "next_q8_output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped", - "symbol": "qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped", - "source": "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_count", - "type": "index" - }, - { - "name": "expert_count", - "type": "index" - }, - { - "name": "output_size", - "type": "index" - } - ], - "bindings": [ - "q8_input", - "expert_table", - "gate_weight", - "up_weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_linked", - "primary_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" - ], - "library_sources": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_linear_q4k_f16_wmma", - "symbol": "qwen3_moe_routed_linear_q4k_f16_wmma", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "input", - "expert_table", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_swiglu_f16", - "symbol": "qwen3_moe_routed_swiglu_f16", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "gate", - "up", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_routed_swiglu_f32", - "symbol": "qwen3_moe_routed_swiglu_f32", - "source": "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - } - ], - "bindings": [ - "gate", - "up", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "primary_sources": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" - ], - "library_sources": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "qwen3_moe_router_projection_f32_four_row_wave32", - "symbol": "qwen3_moe_router_projection_f32_four_row_wave32", - "source": "qwen3_moe/router_projection_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": "archive", - "link_module": "router_projection_f32_linked", - "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen3_moe_router_projection_f32_one_row_wave64", - "symbol": "qwen3_moe_router_projection_f32_one_row_wave64", - "source": "qwen3_moe/router_projection_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": "archive", - "link_module": "router_projection_f32_linked", - "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen3_moe_router_projection_f32_reference", - "symbol": "qwen3_moe_router_projection_f32_reference", - "source": "qwen3_moe/router_projection_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": "archive", - "link_module": "router_projection_f32_linked", - "primary_sources": [ - "qwen3_moe/router_projection_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen3_moe_router_projection_top8_fused_decode_f32", - "symbol": "qwen3_moe_router_projection_top8_fused_decode_f32", - "source": "qwen3_moe/router_projection_top8_fused_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ], - "bindings": [ - "input", - "weight", - "logits", - "completion_counter", - "route_ids", - "route_weights" - ], - "binding_access": [ - "read", - "read", - "read_write", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "router_projection_top8_fused_f32_linked", - "primary_sources": [ - "qwen3_moe/router_projection_top8_fused_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ] - }, - { - "name": "qwen3_moe_router_top8_f32", - "symbol": "qwen3_moe_router_top8_f32", - "source": "qwen3_moe/router_top8_f32.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "route_id_stride", - "type": "index" - } - ], - "bindings": [ - "logits", - "route_ids", - "route_weights" - ], - "binding_access": [ - "read", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "router_top8_f32_linked", - "primary_sources": [ - "qwen3_moe/router_top8_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen3_moe_router_top8_wide_stride_reference", - "symbol": "qwen3_moe_router_top8_wide_stride_reference", - "source": "qwen3_moe/router_top8_f32.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "route_ids" - ], - "binding_access": [ - "read" - ], - "target_selector": "", - "compile_recipe": { - "mode": "archive", - "link_module": "router_top8_f32_linked", - "primary_sources": [ - "qwen3_moe/router_top8_f32.loom" - ], - "library_sources": [ - "qwen3_moe/model_config.loom" - ] - }, - "compile_dependencies": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "qwen_attention_context_base_capture", - "symbol": "qwen_attention_context_base_capture", - "source": "../qwen_owned/attention_state_initialize.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "positions", - "control" - ], - "binding_access": [ - "read", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/attention_state_initialize.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen_attention_decode_state_initialize", - "symbol": "qwen_attention_decode_state_initialize", - "source": "../qwen_owned/attention_state_initialize.loom", - "workload_parameters": [ - { - "name": "context_capacity", - "type": "index" - }, - { - "name": "completion_counter_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "context_capacity", - "type": "index" - }, - { - "name": "completion_counter_count", - "type": "index" - } - ], - "bindings": [ - "positions", - "key_cache_indices", - "value_cache_indices", - "attention_mask", - "completion_counters" - ], - "binding_access": [ - "read", - "write", - "write", - "write", - "write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/attention_state_initialize.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen_attention_metadata", - "symbol": "qwen_attention_metadata", - "source": "../qwen_owned/attention_metadata.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ], - "bindings": [ - "control", - "positions", - "key_cache_indices", - "value_cache_indices", - "attention_mask" - ], - "binding_access": [ - "read", - "read_write", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/attention_metadata.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen_attention_metadata_bringup_workaround", - "symbol": "qwen_attention_metadata_bringup_workaround", - "source": "../qwen_owned/attention_metadata_bringup_workaround.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "context_capacity", - "type": "index" - } - ], - "bindings": [ - "control", - "positions", - "key_cache_indices", - "value_cache_indices", - "attention_mask" - ], - "binding_access": [ - "read", - "read_write", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/attention_metadata_bringup_workaround.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen_decode_attention_metadata", - "symbol": "qwen_decode_attention_metadata", - "source": "../qwen_owned/attention_metadata.loom", - "workload_parameters": [], - "launch_parameters": [], - "bindings": [ - "control", - "positions", - "key_cache_indices", - "value_cache_indices" - ], - "binding_access": [ - "read", - "read_write", - "read_write", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/attention_metadata.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen_token_embedding_q4k", - "symbol": "qwen_token_embedding_q4k", - "source": "../qwen_owned/token_embedding_q4k.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "vocabulary_count", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "vocabulary_count", - "type": "index" - } - ], - "bindings": [ - "token_ids", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/token_embedding_q4k.loom" - ], - "library_sources": [] - }, - "compile_dependencies": [] - }, - { - "name": "qwen_token_embedding_q4k_bringup_workaround", - "symbol": "qwen_token_embedding_q4k_bringup_workaround", - "source": "../qwen_owned/token_embedding_bringup_workaround.loom", - "workload_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "vocabulary_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ], - "launch_parameters": [ - { - "name": "token_count", - "type": "index" - }, - { - "name": "vocabulary_count", - "type": "index" - }, - { - "name": "hidden_size", - "type": "index" - } - ], - "bindings": [ - "token_ids", - "weight", - "output" - ], - "binding_access": [ - "read", - "read", - "read_write" - ], - "target_selector": "", - "compile_recipe": { - "mode": "direct", - "primary_sources": [ - "../qwen_owned/token_embedding_bringup_workaround.loom" + "../qwen/attention_metadata.loom" ], "library_sources": [] - }, - "compile_dependencies": [] - } - ], - "link_modules": [ - { - "name": "router_projection_f32_linked", - "srcs": [ - "qwen3_moe/router_projection_f32.loom" - ], - "libraries": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "router_top8_f32_linked", - "srcs": [ - "qwen3_moe/router_top8_f32.loom" - ], - "libraries": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "router_projection_top8_fused_f32_linked", - "srcs": [ - "qwen3_moe/router_projection_top8_fused_f32.loom" - ], - "libraries": [ - "qwen3_moe/model_config.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_top8_f32.loom" - ] - }, - { - "name": "attention_postprocess_f32_f16_linked", - "srcs": [ - "qwen3_moe/attention_postprocess_f32_f16.loom" - ], - "libraries": [ - "qwen3_moe/model_config.loom" - ] - }, - { - "name": "attention_prepare_quantized_linked", - "srcs": [ - "qwen3_moe/attention_prepare_quantized.loom" - ], - "libraries": [ - "qwen3_moe/model_config.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "attention_qkv_quantized_linked", - "srcs": [ - "qwen3_moe/attention_qkv_quantized.loom" - ], - "libraries": [ - ":dense_linear_quantized_f16_wmma_linked" - ] - }, - { - "name": "attention_qkv_postprocess_fused_linked", - "srcs": [ - "qwen3_moe/attention_qkv_postprocess_fused.loom" - ], - "libraries": [ - ":attention_postprocess_f32_f16_linked", - ":attention_qkv_quantized_linked" - ] - }, - { - "name": "attention_qkv_same_format_prefill_linked", - "srcs": [ - "qwen3_moe/attention_qkv_same_format_prefill.loom" - ], - "libraries": [ - ":dense_linear_quantized_f16_wmma_linked" - ] - }, + }, + "compile_dependencies": [] + } + ], + "link_modules": [ { - "name": "linear_q6k_q8_1_x4_linked", + "name": "router_projection_f32_linked", "srcs": [ - "ggml/linear_q6k_q8_1_x4.loom" + "qwen3_moe/router_projection_f32.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom" ], "libraries": [ - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/model_config.loom" ] }, { - "name": "linear_q6k_f32_linked", + "name": "router_top8_f32_linked", "srcs": [ - "ggml/linear_q6k_f32.loom" + "qwen3_moe/router_top8_f32.loom", + "../loom-libs/motifs/routing/router_topk_row.loom" ], "libraries": [ - ":linear_q6k_q8_1_x4_linked" + "qwen3_moe/model_config.loom" ] }, { - "name": "routed_down_quantized_f16_wmma_linked", + "name": "router_projection_top8_fused_f32_linked", "srcs": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom" + "qwen3_moe/router_projection_top8_fused_f32.loom", + "../loom-libs/motifs/routing/router_projection_storage.loom", + "../loom-libs/motifs/routing/router_topk_row.loom", + "../loom-libs/motifs/routing/router_projection_fused.loom" ], "libraries": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/router_projection_f32.loom", + "qwen3_moe/router_top8_f32.loom" ] }, { - "name": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", + "name": "attention_postprocess_f32_f16_linked", "srcs": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" + "qwen3_moe/attention_postprocess_f32_f16.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ], "libraries": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "qwen3_moe/model_config.loom" ] }, { - "name": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked", + "name": "attention_prepare_quantized_linked", "srcs": [ - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom" + "qwen3_moe/attention_prepare_quantized.loom", + "qwen3_moe/motifs/attention_prepare_target.loom" ], "libraries": [ - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "routed_down_q4k_linked", + "name": "attention_qkv_quantized_linked", "srcs": [ - "qwen3_moe/routed_down_q4k.loom" + "qwen3_moe/attention_qkv_quantized.loom" ], "libraries": [ - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q4_q6_q8_1_x4_qkv.loom", + ":dense_linear_quantized_f16_wmma_linked", + ":attention_prepare_quantized_linked" ] }, { - "name": "routed_down_q6k_linked", + "name": "attention_qkv_postprocess_fused_linked", "srcs": [ - "qwen3_moe/routed_down_q6k.loom" + "qwen3_moe/attention_qkv_postprocess_fused.loom", + "../loom-libs/motifs/attention/rmsnorm_neox_cache_f16.loom" ], "libraries": [ - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_next_q8.loom" + "qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom", + ":attention_qkv_quantized_linked" ] }, { - "name": "routed_gate_up_swiglu_q4k_linked", + "name": "linear_q6k_q8_1_x4_linked", "srcs": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + "ggml/linear_q6k_q8_1_x4.loom" ], "libraries": [ - "ggml/quantize_q8_1_x4.loom" + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { - "name": "batched_decode_expert_dispatch_linked", + "name": "routed_down_quantized_f16_wmma_linked", "srcs": [ - "qwen3_moe/batched_decode_expert_dispatch.loom" + "qwen3_moe/routed_down_quantized_f16_wmma.loom" ], "libraries": [ - "qwen3_moe/model_config.loom" + "../loom-libs/motifs/routing/weighted_reduce_f16_f32.loom" ] }, { - "name": "batched_decode_gate_up_q4k_linked", + "name": "routed_down_weighted_reduce_next_rmsnorm_f32_linked", "srcs": [ - "qwen3_moe/batched_decode_gate_up_q4k.loom" + "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom" ], "libraries": [ + "../loom-libs/motifs/routing/weighted_reduce_next_rmsnorm_f32.loom", + "../loom-libs/motifs/quantize_q8_1_x4.loom", "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/batched_decode_expert_dispatch.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom" + "../loom-libs/motifs/quantize_q8_1_x4_group.loom" ] }, { "name": "expert_table_partition_fused_linked", "srcs": [ - "qwen3_moe/expert_table_partition_fused.loom" - ], - "libraries": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "routed_gate_up_swiglu_q4k_f16_wmma_linked", - "srcs": [ - "qwen3_moe/routed_linear_q4k_f16_wmma.loom" + "qwen3_moe/expert_table_partition_fused.loom", + "qwen3_moe/motifs/expert_table_partition_prefill_512.loom" ], "libraries": [ - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom", + "../loom-libs/ops/matmul_id/vector_pair_binary.loom", + "../loom-libs/motifs/matmul_id/vector_pair.loom", + "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", + "../loom-libs/motifs/elementwise/binary_apply.loom", + "../loom-libs/ops/moe_routing_tables.loom" ] }, { "name": "dense_linear_quantized_f16_wmma_linked", "srcs": [ - "qwen3_moe/dense_linear_quantized_f16_wmma.loom" + "qwen3_moe/dense_linear_quantized_f16_wmma.loom", + "qwen3_moe/motifs/dense_linear_q4k_next_q8.loom", + "qwen3_moe/motifs/attention_prepare_target.loom" ], "libraries": [ + "../loom-libs/motifs/matmul/q4_q8_1_x4_dense.loom", + "../loom-libs/motifs/matmul/q4_q8_1_x4_decode.loom", "ggml/linear_q6k_q8_1_x4.loom", - "qwen3_moe/attention_prepare_quantized.loom", + "../loom-libs/motifs/matmul/q6_q8_1_x4_decode.loom", + "../loom-libs/motifs/q6_k_f16.loom", + "../loom-libs/motifs/normalization/rmsnorm_q8_1_x4.loom", "qwen3_moe/model_config.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "ggml/quantize_q8_1_x4.loom" - ] - }, - { - "name": "flash_attention_decode_split_f32_f16_wmma_next_q8_linked", - "srcs": [ - "qwen3_moe/flash_attention_decode_split_next_q8_test.loom" - ], - "libraries": [ - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "ggml/quantize_q8_1_x4.loom" + "ggml/quantize_q8_1_x4.loom", + "../loom-libs/motifs/quantize_q8_1_x4_group.loom", + "../loom-libs/ops/matmul_id/vector_pair_binary.loom", + "../loom-libs/motifs/matmul_id/vector_pair.loom", + "../loom-libs/motifs/matmul_id/vector_q4k_q8.loom", + "../loom-libs/motifs/elementwise/binary_apply.loom", + "../loom-libs/ops/moe_routing_tables.loom" ] } ], @@ -4771,6 +1180,7 @@ "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", "--config=qwen3_moe.model.rms_epsilon=0.000001", "--config=qwen3_moe.workload.token_capacity=1", "--dry-run", @@ -4811,14 +1221,15 @@ "link_module": "attention_qkv_quantized_linked" }, { - "name": "attention_qkv_postprocess_fused_plan_test", + "name": "attention_qkv_postprocess_fused_q4_differential_plan_test", "args": [ "$(location :attention_qkv_postprocess_fused_linked)", - "--benchmark=@qwen3_moe_attention_qkv_postprocess_fused_boundary_decode", + "--benchmark=@qwen3_moe_attention_qkv_postprocess_fused_differential", "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", - "--config=qwen3_moe.attention.value_uses_q6=1", + "--config=qwen3_moe.attention.rope_mscale=1", + "--config=qwen3_moe.attention.value_uses_q6=0", "--config=qwen3_moe.dense_quantized.input_size=2048", "--config=qwen3_moe.dense_quantized.output_accumulation=0", "--config=qwen3_moe.dense_quantized.output_size=512", @@ -4832,56 +1243,58 @@ "link_module": "attention_qkv_postprocess_fused_linked" }, { - "name": "attention_qkv_same_format_prefill_plan_test", + "name": "attention_qkv_postprocess_fused_q6_differential_plan_test", "args": [ - "$(location :attention_qkv_same_format_prefill_linked)", - "--benchmark=@qwen3_moe_attention_qkv_q4_prefill_512_fused", + "$(location :attention_qkv_postprocess_fused_linked)", + "--benchmark=@qwen3_moe_attention_qkv_postprocess_fused_differential", + "--config=qwen3_moe.attention.head_size=128", "--config=qwen3_moe.attention.key_value_size=512", "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", + "--config=qwen3_moe.attention.value_uses_q6=1", + "--config=qwen3_moe.dense_quantized.input_size=2048", "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.dense_quantized.output_size=512", "--config=qwen3_moe.model.hidden_size=2048", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "attention_qkv_same_format_prefill_linked" - }, - { - "name": "quantize_q8_1_x4_plan_test", - "args": [ - "$(location ggml/quantize_q8_1_x4.loom)", - "--config=ggml.quantize_q8_1_x4.group_capacity=8192", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=ggml.quantize_q8_1_x4.group_capacity=16", + "--config=qwen3_moe.workload.token_capacity=1", "--dry-run", "--output-format=jsonl" ], - "source": "ggml/quantize_q8_1_x4.loom" + "link_module": "attention_qkv_postprocess_fused_linked" }, { - "name": "routed_gate_up_swiglu_q4k_plan_test", + "name": "attention_qkv_postprocess_fused_plan_test", "args": [ - "$(location :routed_gate_up_swiglu_q4k_linked)", - "--benchmark=@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_512", - "--config=qwen3_moe.routed_gate_up.expert_count=128", - "--config=qwen3_moe.routed_gate_up.input_size=2048", - "--config=qwen3_moe.routed_gate_up.output_size=768", - "--config=qwen3_moe.routed_gate_up.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", + "$(location :attention_qkv_postprocess_fused_linked)", + "--benchmark=@qwen3_moe_attention_qkv_postprocess_fused_boundary_decode", + "--config=qwen3_moe.attention.head_size=128", + "--config=qwen3_moe.attention.key_value_size=512", + "--config=qwen3_moe.attention.query_size=4096", + "--config=qwen3_moe.attention.rope_mscale=1", + "--config=qwen3_moe.attention.value_uses_q6=1", + "--config=qwen3_moe.dense_quantized.input_size=2048", + "--config=qwen3_moe.dense_quantized.output_accumulation=0", + "--config=qwen3_moe.dense_quantized.output_size=512", + "--config=qwen3_moe.model.hidden_size=2048", + "--config=qwen3_moe.model.rms_epsilon=0.000001", + "--config=ggml.quantize_q8_1_x4.group_capacity=16", + "--config=qwen3_moe.workload.token_capacity=1", "--dry-run", "--output-format=jsonl" ], - "link_module": "routed_gate_up_swiglu_q4k_linked" + "link_module": "attention_qkv_postprocess_fused_linked" }, { - "name": "linear_q6k_f32_plan_test", + "name": "quantize_q8_1_x4_plan_test", "args": [ - "$(location :linear_q6k_f32_linked)", - "--benchmark=@ggml_linear_q6k_f32_wave64_dense_v_decode", - "--config=ggml.linear_q6k_f32.output_capacity=512", - "--config=ggml.linear_q6k_f32.token_capacity=2048", + "$(location ggml/quantize_q8_1_x4.loom)", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", "--dry-run", "--output-format=jsonl" ], - "link_module": "linear_q6k_f32_linked" + "source": "ggml/quantize_q8_1_x4.loom" }, { "name": "linear_q6k_q8_1_x4_plan_test", @@ -4905,121 +1318,32 @@ "--config=qwen3_moe.routed_down.route_count=8", "--config=qwen3_moe.workload.token_capacity=512", "--dry-run", - "--output-format=jsonl" + "--output-format=jsonl", + "--config=ggml.moe_routing.descriptor_expert_mask=127", + "--config=ggml.moe_routing.descriptor_partition_shift=7", + "--config=ggml.moe_routing.descriptor_row_count_shift=13", + "--config=ggml.moe_routing.expert_count=128", + "--config=ggml.moe_routing.partition_workgroup_size=128", + "--config=ggml.moe_routing.route_count=8", + "--config=ggml.workload.token_capacity=2048" ], "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked" }, { - "name": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_plan_test", + "name": "routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8_plan_test", "args": [ - "$(location :routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked)", - "--benchmark=@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_14", + "$(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=512", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_linked" - }, - { - "name": "routed_down_quantized_f16_wmma_plan_test", - "args": [ - "$(location :routed_down_quantized_f16_wmma_linked)", - "--config=qwen3_moe.routed_down.expert_count=128", - "--config=qwen3_moe.routed_down.input_size=768", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "routed_down_quantized_f16_wmma_linked" - }, - { - "name": "routed_down_q4k_plan_test", - "args": [ - "$(location :routed_down_q4k_linked)", - "--benchmark=@qwen3_moe_routed_down_q4k_q8_1_x4_prefill_512", - "--config=qwen3_moe.routed_down.expert_count=128", - "--config=qwen3_moe.routed_down.input_size=768", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "routed_down_q4k_linked" - }, - { - "name": "routed_down_q6k_plan_test", - "args": [ - "$(location :routed_down_q6k_linked)", - "--benchmark=@qwen3_moe_routed_down_q6k_f32_wave64_decode", - "--config=qwen3_moe.routed_down.expert_count=128", - "--config=qwen3_moe.routed_down.input_size=768", - "--config=qwen3_moe.routed_down.output_size=2048", - "--config=qwen3_moe.routed_down.route_count=8", - "--config=qwen3_moe.workload.token_capacity=1", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "routed_down_q6k_linked" - }, - { - "name": "batched_decode_gate_up_q4k_plan_test", - "args": [ - "$(location :batched_decode_gate_up_q4k_linked)", - "--benchmark=@qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark", - "--config=qwen3_moe.batched_decode.rows2_descriptor_capacity=64", - "--config=qwen3_moe.batched_decode.schedule0_row_limit=1", - "--config=qwen3_moe.batched_decode.schedule1_row_limit=2", - "--config=qwen3_moe.batched_decode.schedule2_row_limit=4", - "--config=qwen3_moe.router.expert_count=128", - "--config=qwen3_moe.router.route_count=8", - "--config=qwen3_moe.routed_gate_up.expert_count=128", - "--config=qwen3_moe.routed_gate_up.input_size=2048", - "--config=qwen3_moe.routed_gate_up.output_size=768", - "--config=qwen3_moe.routed_gate_up.route_count=8", - "--config=qwen3_moe.workload.token_capacity=16", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "batched_decode_gate_up_q4k_linked" - }, - { - "name": "batched_decode_expert_dispatch_plan_test", - "args": [ - "$(location :batched_decode_expert_dispatch_linked)", - "--benchmark=@qwen3_moe_batched_decode_expert_dispatch_diverse", - "--config=qwen3_moe.batched_decode.schedule0_row_limit=1", - "--config=qwen3_moe.batched_decode.schedule1_row_limit=2", - "--config=qwen3_moe.batched_decode.schedule2_row_limit=4", - "--config=qwen3_moe.router.expert_count=128", - "--config=qwen3_moe.router.route_count=8", - "--config=qwen3_moe.workload.token_capacity=16", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "batched_decode_expert_dispatch_linked" - }, - { - "name": "batched_decode_expert_dispatch_configurable_plan_test", - "args": [ - "$(location :batched_decode_expert_dispatch_linked)", - "--benchmark=@qwen3_moe_batched_decode_expert_dispatch_configurable", - "--config=qwen3_moe.batched_decode.schedule0_row_limit=1", - "--config=qwen3_moe.batched_decode.schedule1_row_limit=3", - "--config=qwen3_moe.batched_decode.schedule2_row_limit=6", - "--config=qwen3_moe.router.expert_count=32", - "--config=qwen3_moe.router.route_count=4", - "--config=qwen3_moe.workload.token_capacity=8", + "--config=qwen3_moe.workload.token_capacity=2", + "--config=ggml.quantize_q8_1_x4.group_capacity=32", "--dry-run", "--output-format=jsonl" ], - "link_module": "batched_decode_expert_dispatch_linked" + "link_module": "routed_down_weighted_reduce_next_rmsnorm_f32_linked" }, { "name": "expert_table_partition_fused_plan_test", @@ -5027,38 +1351,17 @@ "$(location :expert_table_partition_fused_linked)", "--benchmark=@qwen3_moe_expert_table_partition_fused_prefill_512", "--dry-run", - "--output-format=jsonl" + "--output-format=jsonl", + "--config=ggml.moe_routing.descriptor_expert_mask=127", + "--config=ggml.moe_routing.descriptor_partition_shift=7", + "--config=ggml.moe_routing.descriptor_row_count_shift=13", + "--config=ggml.moe_routing.expert_count=128", + "--config=ggml.moe_routing.partition_workgroup_size=128", + "--config=ggml.moe_routing.route_count=8", + "--config=ggml.workload.token_capacity=2048" ], "link_module": "expert_table_partition_fused_linked" }, - { - "name": "dense_linear_q6k_f16_wmma_plan_test", - "args": [ - "$(location :dense_linear_quantized_f16_wmma_linked)", - "--benchmark=@qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_128", - "--config=qwen3_moe.dense_quantized.input_size=2048", - "--config=qwen3_moe.dense_quantized.output_accumulation=0", - "--config=qwen3_moe.dense_quantized.output_size=512", - "--config=qwen3_moe.workload.token_capacity=128", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "dense_linear_quantized_f16_wmma_linked" - }, - { - "name": "dense_linear_q4k_f16_wmma_o_plan_test", - "args": [ - "$(location :dense_linear_quantized_f16_wmma_linked)", - "--benchmark=@qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_512", - "--config=qwen3_moe.dense_quantized.input_size=4096", - "--config=qwen3_moe.dense_quantized.output_accumulation=1", - "--config=qwen3_moe.dense_quantized.output_size=2048", - "--config=qwen3_moe.workload.token_capacity=512", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "dense_linear_quantized_f16_wmma_linked" - }, { "name": "dense_linear_q4k_q8_1_x4_plan_test", "args": [ @@ -5068,8 +1371,16 @@ "--config=qwen3_moe.dense_quantized.output_accumulation=0", "--config=qwen3_moe.dense_quantized.output_size=512", "--config=qwen3_moe.workload.token_capacity=1", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", "--dry-run", - "--output-format=jsonl" + "--output-format=jsonl", + "--config=ggml.moe_routing.descriptor_expert_mask=127", + "--config=ggml.moe_routing.descriptor_partition_shift=7", + "--config=ggml.moe_routing.descriptor_row_count_shift=13", + "--config=ggml.moe_routing.expert_count=128", + "--config=ggml.moe_routing.partition_workgroup_size=128", + "--config=ggml.moe_routing.route_count=8", + "--config=ggml.workload.token_capacity=2048" ], "link_module": "dense_linear_quantized_f16_wmma_linked" }, @@ -5084,140 +1395,28 @@ "--config=qwen3_moe.model.hidden_size=2048", "--config=qwen3_moe.model.rms_epsilon=0.000001", "--config=qwen3_moe.workload.token_capacity=1", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "dense_linear_quantized_f16_wmma_linked" - }, - { - "name": "routed_gate_up_swiglu_q4k_f16_wmma_plan_test", - "args": [ - "$(location :routed_gate_up_swiglu_q4k_f16_wmma_linked)", - "--config=qwen3_moe.routed_gate_up.expert_count=128", - "--config=qwen3_moe.routed_gate_up.input_size=2048", - "--config=qwen3_moe.routed_gate_up.output_size=768", - "--config=qwen3_moe.routed_gate_up.route_count=8", - "--config=qwen3_moe.workload.token_capacity=2048", - "--dry-run", - "--output-format=jsonl" - ], - "link_module": "routed_gate_up_swiglu_q4k_f16_wmma_linked" - }, - { - "name": "flash_attention_decode_f32_f16_wmma_plan_test", - "args": [ - "$(location qwen3_moe/flash_attention_decode_f32_f16_wmma.loom)", - "--config=qwen3_moe.attention.decode.output_partition_count=1", - "--config=qwen3_moe.attention.key_value_head_count=4", - "--config=qwen3_moe.attention.query_head_count=32", - "--dry-run", - "--output-format=jsonl" - ], - "source": "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom" - }, - { - "name": "flash_attention_decode_q128_f32_f16_wmma_plan_test", - "args": [ - "$(location qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom)", - "--config=qwen3_moe.attention.key_value_head_count=4", - "--config=qwen3_moe.attention.query_head_count=32", - "--dry-run", - "--output-format=jsonl" - ], - "source": "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom" - }, - { - "name": "flash_attention_decode_split_f32_f16_wmma_plan_test", - "args": [ - "$(location qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom)", - "--config=qwen3_moe.attention.key_value_head_count=4", - "--config=qwen3_moe.attention.key_value_token_capacity=2048", - "--config=qwen3_moe.attention.query_head_count=32", - "--dry-run", - "--output-format=jsonl" - ], - "source": "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom" - }, - { - "name": "flash_attention_f32_f16_wmma_plan_test", - "args": [ - "$(location qwen3_moe/flash_attention_f32_f16_wmma.loom)", - "--config=qwen3_moe.attention.key_value_head_count=4", - "--config=qwen3_moe.attention.query_head_count=32", - "--config=qwen3_moe.workload.token_capacity=2048", - "--dry-run", - "--output-format=jsonl" - ], - "source": "qwen3_moe/flash_attention_f32_f16_wmma.loom" - }, - { - "name": "owned_token_embedding_decode_plan_test", - "args": [ - "$(location ../qwen_owned/token_embedding_q4k.loom)", - "--benchmark=@qwen_token_embedding_q4k_decode", - "--dry-run", - "--output-format=jsonl", - "--sample-compilation=per_sample" - ], - "source": "../qwen_owned/token_embedding_q4k.loom" - }, - { - "name": "owned_token_embedding_prefill_plan_test", - "args": [ - "$(location ../qwen_owned/token_embedding_q4k.loom)", - "--benchmark=@qwen_token_embedding_q4k_prefill_512", - "--dry-run", - "--output-format=jsonl", - "--sample-compilation=per_sample" - ], - "source": "../qwen_owned/token_embedding_q4k.loom" - }, - { - "name": "owned_attention_context_base_capture_plan_test", - "args": [ - "$(location ../qwen_owned/attention_state_initialize.loom)", - "--benchmark=@qwen_attention_context_base_capture_benchmark", - "--dry-run", - "--output-format=jsonl", - "--sample-compilation=per_sample" - ], - "source": "../qwen_owned/attention_state_initialize.loom", - "owner": "ggml-hrx" - }, - { - "name": "owned_attention_decode_state_initialize_plan_test", - "args": [ - "$(location ../qwen_owned/attention_state_initialize.loom)", - "--benchmark=@qwen_attention_decode_state_initialize_benchmark", + "--config=ggml.quantize_q8_1_x4.group_capacity=8192", "--dry-run", "--output-format=jsonl", - "--sample-compilation=per_sample" + "--config=ggml.moe_routing.descriptor_expert_mask=127", + "--config=ggml.moe_routing.descriptor_partition_shift=7", + "--config=ggml.moe_routing.descriptor_row_count_shift=13", + "--config=ggml.moe_routing.expert_count=128", + "--config=ggml.moe_routing.partition_workgroup_size=128", + "--config=ggml.moe_routing.route_count=8", + "--config=ggml.workload.token_capacity=2048" ], - "source": "../qwen_owned/attention_state_initialize.loom", - "owner": "ggml-hrx" + "link_module": "dense_linear_quantized_f16_wmma_linked" }, { "name": "owned_attention_metadata_prefill_plan_test", "args": [ - "$(location ../qwen_owned/attention_metadata.loom)", + "$(location ../qwen/attention_metadata.loom)", "--benchmark=@qwen_attention_metadata_prefill_512", "--dry-run", - "--output-format=jsonl", - "--sample-compilation=per_sample" - ], - "source": "../qwen_owned/attention_metadata.loom" - }, - { - "name": "owned_gather_add_plan_test", - "args": [ - "$(location ../hrx_owned/gather_add_f32.loom)", - "--benchmark=@ggml_gather_add_noncontiguous", - "--dry-run", - "--output-format=jsonl", - "--sample-compilation=per_sample" + "--output-format=jsonl" ], - "source": "../hrx_owned/gather_add_f32.loom", - "owner": "ggml-hrx" + "source": "../qwen/attention_metadata.loom" } ] } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom index 43551b88381f..6c97509cdef1 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_postprocess_f32_f16.loom @@ -13,6 +13,8 @@ // distinct so the kernel does not impose an ordering contract on the cache. // The stage scheduler owns those indices and guarantees they select valid cache // rows; violating that trusted contract is an error rather than a skipped write. +template.decl @qwen3_moe.attention.postprocess.head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) + amdgpu.target @qwen3_moe_attention_postprocess_gfx11_wave32 {subgroup_size = 32} config.decl @qwen3_moe.model.rms_epsilon : f32 @@ -23,160 +25,7 @@ config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] -// Applies one RMSNorm and two adjacent NEOX pairs. NEOX pairs corresponding -// channels from the low and high halves, while keeping each half contiguous. -// A two-pair packet therefore gives both halves naturally coalesced loads and -// stores without changing the model's pairing rule. -func.def inline @qwen3_moe_rmsnorm_neox_packet(%row_sum: f32, %head_size: f32, %epsilon: f32, %position: f32, %inverse_frequencies: vector<2xf32>, %low_values: vector<2xf32>, %high_values: vector<2xf32>, %low_weights: vector<2xf32>, %high_weights: vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) { - %mean = scalar.divf %row_sum, %head_size : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - %scale_vector = vector.splat %scale : vector<2xf32> - %position_vector = vector.splat %position : vector<2xf32> - %inverse_two_pi = scalar.constant 0.15915494309189535 : f32 - %inverse_two_pi_vector = vector.splat %inverse_two_pi : vector<2xf32> - %normalized_low = vector.mulf %low_values, %scale_vector : vector<2xf32> - %normalized_high = vector.mulf %high_values, %scale_vector : vector<2xf32> - %scaled_low = vector.mulf %normalized_low, %low_weights : vector<2xf32> - %scaled_high = vector.mulf %normalized_high, %high_weights : vector<2xf32> - %angles = vector.mulf %position_vector, %inverse_frequencies : vector<2xf32> - %turns = vector.mulf %angles, %inverse_two_pi_vector : vector<2xf32> - %cosines = vector.costurnsf %turns : vector<2xf32> - %sines = vector.sinturnsf %turns : vector<2xf32> - %low_cosines = vector.mulf %scaled_low, %cosines : vector<2xf32> - %high_sines = vector.mulf %scaled_high, %sines : vector<2xf32> - %low_sines = vector.mulf %scaled_low, %sines : vector<2xf32> - %high_cosines = vector.mulf %scaled_high, %cosines : vector<2xf32> - %rotated_low = vector.subf %low_cosines, %high_sines : vector<2xf32> - %rotated_high = vector.addf %low_sines, %high_cosines : vector<2xf32> - func.return %rotated_low, %rotated_high : vector<2xf32>, vector<2xf32> -} - -// Processes one head domain after its raw projection row is visible. Callers -// may use a larger workgroup than the packet count; inactive workitems -// contribute zero to normalization and perform no loads or stores. -func.template device @qwen3_moe_attention_postprocess_head_body(%publish_output: i1, %token_count: index, %cache_row_count: index, %head_domain0: index, %token0: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) { - %query_size0 = config.get @qwen3_moe.attention.query_size : index - %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index - %head_size0 = config.get @qwen3_moe.attention.head_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_cache_row_count = index.assume %cache_row_count [range(%cache_row_count, 1, 1048576)] : index - %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), range(%head_size0, 4, 1024), mul(%head_size0, 4), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index - %channel0 = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %half_head_size = index.div %head_size, %c2 : index - %pair_packet_count = index.div %head_size, %c4 : index - %query_head_count = index.div %query_size, %head_size : index - %key_value_head_count = index.div %key_value_size, %head_size : index - %key_value_domain_count = index.mul %key_value_head_count, %c2 : index - %head_domain_count = index.add %query_head_count, %key_value_domain_count : index - %head_domain = index.assume %head_domain0 [lt(%head_domain0, %head_domain_count)] : index - %key_domain_end = index.add %query_head_count, %key_value_head_count : index - %is_query = index.cmp ult, %head_domain, %query_head_count : index - %is_query_or_key = index.cmp ult, %head_domain, %key_domain_end : index - %key_value_head = index.rem %head_domain, %key_value_head_count : index - %active_channel = index.cmp ult, %channel0, %pair_packet_count : index - %head_size_i32 = index.cast %head_size : index to i32 - %head_size_f32 = scalar.sitofp %head_size_i32 : i32 to f32 - %positions_noalias, %key_cache_indices_noalias, %value_cache_indices_noalias, %query_input_noalias, %key_input_noalias, %value_input_noalias, %query_norm_weight_noalias, %key_norm_weight_noalias, %inverse_frequencies_noalias, %query_output_noalias, %key_cache_noalias, %value_cache_noalias = buffer.assume.noalias %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache : buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer - %positions_view = buffer.view %positions_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> - %key_cache_indices_view = buffer.view %key_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> - %value_cache_indices_view = buffer.view %value_cache_indices_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi64> - %query_input_view = buffer.view %query_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - %key_input_view = buffer.view %key_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> - %value_input_view = buffer.view %value_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> - %query_norm_weight_view = buffer.view %query_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> - %key_norm_weight_view = buffer.view %key_norm_weight_noalias[%c0_offset] : buffer -> view<[%head_size]xf32> - %inverse_frequencies_view = buffer.view %inverse_frequencies_noalias[%c0_offset] : buffer -> view<[%half_head_size]xf32> - %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - %key_cache_words_view = buffer.view %key_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - %value_cache_words_view = buffer.view %value_cache_noalias[%c0_offset] : buffer -> view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - %row_sum = scf.if %is_query_or_key -> (f32) { - %partial_sum = scf.if %publish_output -> (f32) { - %channel_sum = scf.if %active_channel -> (f32) { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %reduction_channel = index.mul %channel, %c4 : index - %reduction_values = scf.if %is_query -> (vector<4xf32>) { - %query_values = vector.load %query_input_view[%token, %head_domain, %reduction_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<4xf32> - scf.yield %query_values : vector<4xf32> - } else { - %key_values = vector.load %key_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> - scf.yield %key_values : vector<4xf32> - } - %squares = vector.mulf %reduction_values, %reduction_values : vector<4xf32> - %sum = vector.reduce %squares, %c0_f32 : vector<4xf32>, f32 - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.yield %channel_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %sum = kernel.workgroup.reduce %partial_sum : f32 - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.if %publish_output { - scf.if %active_channel { - %channel = index.assume %channel0 [lt(%channel0, %pair_packet_count)] : index - %pair_channel = index.mul %channel, %c2 : index - %paired_channel = index.add %pair_channel, %half_head_size : index - %paired_word = index.add %channel, %pair_packet_count : index - %reduction_channel = index.mul %channel, %c4 : index - scf.if %is_query { - %low_values = vector.load %query_input_view[%token, %head_domain, %pair_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> - %high_values = vector.load %query_input_view[%token, %head_domain, %paired_channel] : view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> -> vector<2xf32> - %low_weights = vector.load %query_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> - %high_weights = vector.load %query_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> - %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> - %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - vector.store %rotated_low, %query_output_view[%token, %head_domain, %pair_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - vector.store %rotated_high, %query_output_view[%token, %head_domain, %paired_channel] : vector<2xf32>, view<[%launch_token_count]x[%query_head_count]x[%head_size]xf32> - } else { - scf.if %is_query_or_key { - %low_values = vector.load %key_input_view[%token, %key_value_head, %pair_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> - %high_values = vector.load %key_input_view[%token, %key_value_head, %paired_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<2xf32> - %low_weights = vector.load %key_norm_weight_view[%pair_channel] : view<[%head_size]xf32> -> vector<2xf32> - %high_weights = vector.load %key_norm_weight_view[%paired_channel] : view<[%head_size]xf32> -> vector<2xf32> - %inverse_frequencies_packet = vector.load %inverse_frequencies_view[%pair_channel] : view<[%half_head_size]xf32> -> vector<2xf32> - %position_i32 = view.load %positions_view[%token] : view<[%launch_token_count]xi32> -> i32 - %position = scalar.sitofp %position_i32 : i32 to f32 - %cache_index_raw = view.load %key_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 - %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 - %cache_index0 = index.cast %cache_index_i64 : i64 to index - %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index - %rotated_low, %rotated_high = func.call @qwen3_moe_rmsnorm_neox_packet(%row_sum, %head_size_f32, %epsilon, %position, %inverse_frequencies_packet, %low_values, %high_values, %low_weights, %high_weights) : (f32, f32, f32, f32, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>, vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) - %half_low = vector.fptrunc %rotated_low : vector<2xf32> to vector<2xf16> - %half_high = vector.fptrunc %rotated_high : vector<2xf32> to vector<2xf16> - %packed_low = vector.bitcast %half_low : vector<2xf16> to vector<1xi32> - %packed_high = vector.bitcast %half_high : vector<2xf16> to vector<1xi32> - vector.store %packed_low, %key_cache_words_view[%cache_index, %key_value_head, %channel] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - vector.store %packed_high, %key_cache_words_view[%cache_index, %key_value_head, %paired_word] : vector<1xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - } else { - %cache_index_raw = view.load %value_cache_indices_view[%token] : view<[%launch_token_count]xi64> -> i64 - %cache_index_i64 = scalar.assume %cache_index_raw [range(%cache_index_raw, 0, 1048575)] : i64 - %cache_index0 = index.cast %cache_index_i64 : i64 to index - %cache_index = index.assume %cache_index0 [lt(%cache_index0, %bounded_cache_row_count)] : index - %values = vector.load %value_input_view[%token, %key_value_head, %reduction_channel] : view<[%launch_token_count]x[%key_value_head_count]x[%head_size]xf32> -> vector<4xf32> - %half_values = vector.fptrunc %values : vector<4xf32> to vector<4xf16> - %packed_values = vector.bitcast %half_values : vector<4xf16> to vector<2xi32> - vector.store %packed_values, %value_cache_words_view[%cache_index, %key_value_head, %pair_channel] : vector<2xi32>, view<[%bounded_cache_row_count]x[%key_value_head_count]x[%half_head_size]xi32> - } - } - } - } - func.return -} +config.decl @qwen3_moe.attention.rope_mscale : f32 kernel.def target(@qwen3_moe_attention_postprocess_gfx11_wave32) @qwen3_moe_attention_postprocess_f32_f16(%token_count: index, %cache_row_count: index) { %query_size = config.get @qwen3_moe.attention.query_size : index @@ -197,7 +46,7 @@ kernel.def target(@qwen3_moe_attention_postprocess_gfx11_wave32) @qwen3_moe_atte %c0 = index.constant 0 : index %valid_token = index.cmp ult, %token0, %token_count : index %safe_token = scf.select %valid_token, %token0, %c0 : index - func.apply(%valid_token, %token_count, %cache_row_count, %head_domain, %safe_token, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe.attention.postprocess.head_body>(%valid_token, %token_count, %cache_row_count, %head_domain, %safe_token, %positions, %key_cache_indices, %value_cache_indices, %query_input, %key_input, %value_input, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom index 2b6b45ab8318..c0091da9a755 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_prepare_quantized.loom @@ -15,14 +15,16 @@ // Every workitem packs one four-value word per stripe; adjacent eight-lane // cohorts reduce the maximum and quantized sum for one logical Q8_1 block. // Scratch is fixed at 1152 bytes regardless of hidden size. -amdgpu.target @qwen3_moe_attention_prepare_gfx11_wave32 {subgroup_size = 32} +template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) + +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 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 // Shared Q8_1 x4 packer used by the standalone differential path. -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$10: index, %input_size$11: index) launch(%token_count$12: index, %input_size$13: index, %input: buffer, %output: buffer) // Materializes the ordinary Qwen RMSNorm boundary. This remains useful for // prefill schedules that reuse one normalized row across many output tiles; @@ -101,180 +103,6 @@ kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_rmsnorm_ kernel.return } -// Shared RMSNorm and GGML Q8_1 x4 row producer. The caller owns launch geometry -// and supplies the logical token whose complete hidden row this workgroup owns. -// The attention export discards the ordinary F32 row, while feed-forward -// publishes it for the router without rereading and requantizing the normalized -// values. -func.template device @qwen3_moe_rmsnorm_quantize_q8_1_x4_body(%publish_normalized: i1, %reduction_subgroup_count0: index, %token_count: index, %token0: index, %input: buffer, %weight: buffer, %normalized_output: buffer, %q8_output: buffer) { - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index - %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c1024 = index.constant 1024 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 1024 : offset - %scratch_bytes = index.constant 1152 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %hidden_size_i32 = index.cast %hidden_size : index to i32 - %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 - %physical_group_count = index.div %hidden_size, %c128 : index - %row_bytes = index.scale %physical_group_count, %group_bytes : index, offset -> offset - %token_output_byte_base = index.scale %token, %row_bytes : index, offset -> offset - %input_noalias, %weight_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %q8_output : buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %normalized_output_view = buffer.view %normalized_output[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<256xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<32xf32> - // Reduce the complete row before any block-local quantization. - %thread_sum = scf.for %channel = [%workitem to %hidden_size step %c256](%running_sum = %c0_f32 : f32) -> (f32) { - %value = view.load %input_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> f32 - %square = scalar.mulf %value, %value : f32 - %next_sum = scalar.addf %running_sum, %square : f32 - scf.yield %next_sum : f32 - } - %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<256xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index - %is_reduction_lane = index.cmp ult, %lane, %reduction_subgroup_count : index - %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 - %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { - %value = view.load %scratch_values[%lane] : view<256xf32> -> f32 - scf.yield %value : f32 - } else { - scf.yield %c0_f32 : f32 - } - %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 - %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 - scf.if %writes_scale { - %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_values[%c0] : f32, view<256xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %row_scale = view.load %scratch_values[%c0] : view<256xf32> -> f32 - %row_scale_vector = vector.splat %row_scale : vector<4xf32> - %publishes_normalized = scalar.andi %publish_normalized, %valid_token : i1 - // Reuse one scratch frame for each 1024-element stripe. Hidden sizes need - // only be divisible by 128; inactive workitems in the final stripe carry - // zeros and never publish. - scf.for %stripe_base = [%c0 to %hidden_size step %c1024] { - %word_element_add = index.mul %workitem, %c4 : index - %channel = index.add %stripe_base, %word_element_add : index - %valid_word = index.cmp ult, %channel, %hidden_size : index - %mask = vector.mask.range [%channel to %hidden_size step %c1] : index -> vector<4xi1> - %input_values = vector.load.mask %input_view[%token, %channel], %mask, %c0_f32x4 : view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> - %learned_weights = vector.load.mask %weight_view[%channel], %mask, %c0_f32x4 : view<[%hidden_size]xf32>, vector<4xi1>, vector<4xf32> - %normalized0 = vector.mulf %input_values, %row_scale_vector : vector<4xf32> - %normalized = vector.mulf %normalized0, %learned_weights : vector<4xf32> - scf.if %publishes_normalized { - vector.store.mask %normalized, %normalized_output_view[%token, %channel], %mask : vector<4xf32>, view<[%launch_token_count]x[%hidden_size]xf32>, vector<4xi1> - } - %absolute_values = vector.absf %normalized : vector<4xf32> - %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %word_in_block = index.rem %workitem, %c8 : index - %block_in_stripe = index.div %workitem, %c8 : index - %is_block_leader = index.cmp eq, %word_in_block, %c0 : index - %writes_block_d = scalar.andi %valid_word, %is_block_leader : i1 - scf.if %writes_block_d { - %cohort_base = index.mul %block_in_stripe, %c8 : index - %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> - %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%block_in_stripe] : f32, view<32xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %d = scf.if %valid_word -> (f32) { - %block_d = view.load %scratch_d[%block_in_stripe] : view<32xf32> -> f32 - scf.yield %block_d : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %publishes_q8_word = scalar.andi %valid_word, %valid_token : i1 - scf.if %publishes_q8_word { - %q8_block = index.div %channel, %c32 : index - %physical_group = index.div %q8_block, %c4 : index - %block_in_group = index.rem %q8_block, %c4 : index - %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset - %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %q8_output_noalias[%payload_byte_offset] : buffer -> view<32xi32> - %block_word_base = index.mul %block_in_group, %c8 : index - %packed_word_index0 = index.add %block_word_base, %word_in_block : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> - } - %thread_quantized_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_quantized_sum, %scratch_values[%workitem] : f32, view<256xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %publishes_block_ds = scalar.andi %writes_block_d, %valid_token : i1 - scf.if %publishes_block_ds { - %cohort_base = index.mul %block_in_stripe, %c8 : index - %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %q8_block = index.div %channel, %c32 : index - %physical_group = index.div %q8_block, %c4 : index - %block_in_group = index.rem %q8_block, %c4 : index - %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset - %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset - %group_ds = buffer.view %q8_output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - %s_index = index.add %ds_index, %c1 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - } - func.return -} - // Fuses attention RMSNorm directly into the Q8_1 x4 row consumed by decode // projections. The output has hidden_size / 128 physical groups of 144 bytes. kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) { @@ -286,7 +114,7 @@ kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_attentio %reduction_subgroup_count = index.constant 8 : index %token = kernel.workgroup.id : index %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer - func.apply(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input_noalias, %weight_noalias, %output_noalias, %output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input_noalias, %weight_noalias, %output_noalias, %output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) kernel.return } @@ -301,7 +129,7 @@ kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_rmsnorm_ %reduction_subgroup_count = index.constant 8 : index %token = kernel.workgroup.id : index %input_noalias, %weight_noalias, %normalized_output_noalias, %q8_output_noalias = buffer.assume.noalias %input, %weight, %normalized_output, %q8_output : buffer, buffer, buffer, buffer - func.apply(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input_noalias, %weight_noalias, %normalized_output_noalias, %q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input_noalias, %weight_noalias, %normalized_output_noalias, %q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) kernel.return } @@ -335,7 +163,7 @@ check.case public @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_differential_cas check.case public @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4_multistripe_case { %token_count = check.literal value(14) : index %hidden_size = check.literal value(2048) : index - %input_seed = check.param.seed base(0x514d4f4550524550) count(1) : i64 + %input_seed = check.param.seed base(5858425849397790032) count(1) : i64 %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<14x2048xf32> %weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> %normalized = check.generate.fill value(0.0) : tensor<14x2048xf32> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom index 70b2d9ec490a..62d91587f23f 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_postprocess_fused.loom @@ -12,6 +12,14 @@ // lets the postprocess consume freshly published rows; it does not impose a // hidden cache-index or position relationship. Counters return to zero only // after semantic outputs are visible, preserving reusable command buffers. +template.decl @qwen3_moe.attention.postprocess.head_body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: buffer, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer) + +template.decl @qwen3_moe.attention.qkv_postprocess_fused_decode.body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) + +template.decl @qwen3_moe.attention.qkv_postprocess_fused_decode.launch() -> (index, index, index) + +template.decl @qwen3_moe.attention.qkv_quantized.body(%arg0: index, %arg1: i1, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: buffer, %arg9: buffer, %arg10: buffer) + amdgpu.target @qwen3_moe_attention_qkv_postprocess_gfx11_wave32 {subgroup_size = 32} config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 4, 1024), mul(%value, 4)] @@ -23,73 +31,13 @@ config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%va config.decl @qwen3_moe.attention.value_uses_q6 : %value: index where [range(%value, 0, 1)] // Reference entry points used by differential and benchmark cases. -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.decl @qwen3_moe_attention_qkv_quantized(%token_count: index) launch(%token_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) - -kernel.decl @qwen3_moe_attention_postprocess_f32_f16(%token_count: index, %cache_row_count: index) launch(%token_count: index, %cache_row_count: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) - -// Storage-specific exports pass a constant value format into this shared body, -// while the compatibility export continues to source it from model config. -func.template device @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { - %publish_projection_output = scalar.constant true : i1 - %projection_token = kernel.workgroup.id : index - func.apply(%value_uses_q6_index, %publish_projection_output, %token_count, %projection_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output_raw, %key_output_raw, %value_output_raw) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - %query_size0 = config.get @qwen3_moe.attention.query_size : index - %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index - %head_size0 = config.get @qwen3_moe.attention.head_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 128, 262144), range(%key_value_size0, 128, 262144), range(%head_size0, 128, 128), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c8 = index.constant 8 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %channel_tile = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index - %head_tile_count = index.div %head_size, %c8 : index - %head_domain0 = index.div %channel_tile, %head_tile_count : index - %query_head_count = index.div %query_size, %head_size : index - %key_value_head_count = index.div %key_value_size, %head_size : index - %key_value_domain_count = index.mul %key_value_head_count, %c2 : index - %head_domain_count = index.add %query_head_count, %key_value_domain_count : index - %head_domain, %completion_counter_count = index.assume %head_domain0, %head_domain_count [lt(%head_domain0, %head_domain_count)] : index, index - %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index - %completion_counters_aligned = buffer.assume.alignment %completion_counters {minimum_alignment = 16} : buffer - %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%completion_counter_count]xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - // Publish every producer's projection stores before the leader advances one - // workgroup arrival. The last arrival then acquires the complete head. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - %old_counter = view.atomic.rmw %c1_i32, %completion_counters_view[%head_domain] {ordering = acq_rel, scope = device} : i32, view<[%completion_counter_count]xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %head_tile_count_i32 = index.cast %head_tile_count : index to i32 - %last_head_tile_i32 = scalar.subi %head_tile_count_i32, %c1_i32 : i32 - %negative_head_tile_count_i32 = scalar.subi %c0_i32, %head_tile_count_i32 : i32 - %is_last_head_tile = scalar.cmpi eq, %old_counter, %last_head_tile_i32 : i32 - scf.if %is_last_head_tile { - kernel.barrier scope(workgroup) ordering(acquire) - %publish_postprocess_output = scalar.constant true : i1 - func.apply(%publish_postprocess_output, %bounded_token_count, %cache_row_count, %head_domain, %token, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - // The counter becomes reusable only after all semantic output stores. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - view.atomic.reduce %negative_head_tile_count_i32, %completion_counters_view[%head_domain] {ordering = release, scope = device} : i32, view<[%completion_counter_count]xi32> - } - } - func.return -} +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$55: index, %input_size$56: index) launch(%token_count$57: index, %input_size$58: index, %input: buffer, %output: buffer) -func.template @qwen3_moe_attention_qkv_postprocess_fused_decode_launch() -> (index, index, index) { +kernel.decl @qwen3_moe_attention_qkv_quantized(%token_count$61: index) launch(%token_count$62: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + +kernel.decl @qwen3_moe_attention_postprocess_f32_f16(%token_count$70: index, %cache_row_count$71: index) launch(%token_count$72: index, %cache_row_count$73: index, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_input: buffer, %key_input: buffer, %value_input: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer) + +template.def<@qwen3_moe.attention.qkv_postprocess_fused_decode.launch> @qwen3_moe_attention_qkv_postprocess_fused_decode_launch() -> (index, index, index) { %query_size = config.get @qwen3_moe.attention.query_size : index %key_value_size = config.get @qwen3_moe.attention.key_value_size : index %c1 = index.constant 1 : index @@ -101,33 +49,25 @@ func.template @qwen3_mo %output_size = index.add %query_size, %key_value_output_size : index %padded_output_size = index.add %output_size, %c7 : index %output_tiles = index.div %padded_output_size, %c8 : index - func.return %output_tiles, %c1, %c256 : index, index, index + template.return %output_tiles, %c1, %c256 : index, index, index } kernel.def target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode(%token_count: index, %cache_row_count: index) { - %output_tiles, %c1, %c256 = func.apply() : () -> (index, index, index) + %query_size = config.get @qwen3_moe.attention.query_size : index + %c1 = index.constant 1 : index + %key_value_size = config.get @qwen3_moe.attention.key_value_size : index + %c2 = index.constant 2 : index + %c7 = index.constant 7 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %key_value_output_size = index.mul %key_value_size, %c2 : index + %output_size = index.add %query_size, %key_value_output_size : index + %padded_output_size = index.add %output_size, %c7 : index + %output_tiles = index.div %padded_output_size, %c8 : index kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index } launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { %value_uses_q6_index = config.get @qwen3_moe.attention.value_uses_q6 : index - func.apply(%value_uses_q6_index, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode_q4(%token_count: index, %cache_row_count: index) { - %output_tiles, %c1, %c256 = func.apply() : () -> (index, index, index) - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { - %q4 = index.constant 0 : index - func.apply(%q4, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_attention_qkv_postprocess_gfx11_wave32) @qwen3_moe_attention_qkv_postprocess_fused_decode_q6(%token_count: index, %cache_row_count: index) { - %output_tiles, %c1, %c256 = func.apply() : () -> (index, index, index) - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { - %q6 = index.constant 1 : index - func.apply(%q6, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe.attention.qkv_postprocess_fused_decode.body>(%value_uses_q6_index, %token_count, %cache_row_count, %q8_input, %query_weight, %key_weight, %value_weight, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache, %completion_counters) : (index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } @@ -239,6 +179,8 @@ check.case public @qwen3_moe_attention_qkv_postprocess_fused_benchmark_case { check.return } +check.benchmark<@qwen3_moe_attention_qkv_postprocess_fused_differential_case> @qwen3_moe_attention_qkv_postprocess_fused_differential + check.benchmark<@qwen3_moe_attention_qkv_postprocess_composed_benchmark_case> @qwen3_moe_attention_qkv_postprocess_composed_decode check.benchmark<@qwen3_moe_attention_qkv_postprocess_fused_benchmark_case> @qwen3_moe_attention_qkv_postprocess_fused_boundary_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom index 2cb2ccf5e809..4967c5942e92 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_quantized.loom @@ -10,6 +10,8 @@ // F32 projections. The adjacent preparation producer owns RMSNorm and Q8_1 // packing; the following attention postprocess owns per-head normalization, // RoPE, and K/V cache publication. +template.decl @qwen3_moe.attention.qkv_quantized.body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + amdgpu.target @qwen3_moe_attention_qkv_gfx11_wave32 {subgroup_size = 32} config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] @@ -31,113 +33,14 @@ func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: bu func.decl @ggml_q6k_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) // Fused RMSNorm and Q8_1 x4 producer linked from attention preparation. -kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %q8_output: buffer) +kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count$30: index) launch(%token_count$31: index, %input: buffer, %weight: buffer, %q8_output: buffer) // Reference entry points used only by differential cases. -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.decl @qwen3_moe_dense_linear_q4k_q8_1_x4(%token_count: index) launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer) - -kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) - -// Device body shared by the ordinary projection and completion-fused -// postprocess exports. Each caller owns its boundary after the raw row stores. -func.template device @qwen3_moe_attention_qkv_quantized_body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) { - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %query_size0 = config.get @qwen3_moe.attention.query_size : index - %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %hidden_size = index.assume %hidden_size0 [range(%hidden_size0, 128, 32768), mul(%hidden_size0, 128)] : index - %query_size, %key_value_size = index.assume %query_size0, %key_value_size0 [range(%query_size0, 1, 262144), range(%key_value_size0, 1, 262144), mul(%query_size0, %key_value_size0)] : index, index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144_bytes = index.constant 144 : offset - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %global_channel = index.add %channel_base, %subgroup : index - %key_value_end = index.add %query_size, %key_value_size : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %total_output_size = index.add %query_size, %key_value_output_size : index - %valid_channel = index.cmp ult, %global_channel, %total_output_size : index - %is_query = index.cmp ult, %global_channel, %query_size : index - %is_key = index.cmp ult, %global_channel, %key_value_end : index - %key_value_channel = index.rem %global_channel, %key_value_size : index - %value_uses_q6 = index.cmp eq, %value_uses_q6_index, %c1 : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %quant_block_count = index.div %hidden_size, %c256 : index - %q4_row_bytes = index.scale %quant_block_count, %c144_bytes : index, offset -> offset - %q6_row_bytes = index.scale %quant_block_count, %c210_bytes : index, offset -> offset - %q8_group_count = index.div %hidden_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset - %q8_noalias, %query_weight_noalias, %key_weight_noalias, %value_weight_noalias, %query_output_noalias, %key_output_noalias, %value_output_noalias = buffer.assume.noalias %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer - %lane_sum = scf.if %publish_output -> (f32) { - %channel_sum = scf.if %valid_channel -> (f32) { - %projection_sum = scf.if %is_query -> (f32) { - %row_byte_base = index.scale %global_channel, %q4_row_bytes : index, offset -> offset - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %query_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - %key_or_value_sum = scf.if %is_key -> (f32) { - %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %key_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - %value_sum = scf.if %value_uses_q6 -> (f32) { - %row_byte_base = index.scale %key_value_channel, %q6_row_bytes : index, offset -> offset - %sum = func.call @ggml_q6k_q8_1_x4_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - %row_byte_base = index.scale %key_value_channel, %q4_row_bytes : index, offset -> offset - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%hidden_size, %value_weight_noalias, %row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } - scf.yield %value_sum : f32 - } - scf.yield %key_or_value_sum : f32 - } - scf.yield %projection_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.yield %channel_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %dot = kernel.subgroup.reduce %lane_sum : f32 - %query_output_view = buffer.view %query_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%query_size]xf32> - %key_output_view = buffer.view %key_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> - %value_output_view = buffer.view %value_output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%key_value_size]xf32> - scf.if %publish_output { - scf.if %valid_channel { - scf.if %is_lane_zero { - scf.if %is_query { - view.store %dot, %query_output_view[%token, %global_channel] : f32, view<[%launch_token_count]x[%query_size]xf32> - } else { - scf.if %is_key { - view.store %dot, %key_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> - } else { - view.store %dot, %value_output_view[%token, %key_value_channel] : f32, view<[%launch_token_count]x[%key_value_size]xf32> - } - } - } - } - } - func.return -} +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$35: index, %input_size$36: index) launch(%token_count$37: index, %input_size$38: index, %input: buffer, %output: buffer) + +kernel.decl @qwen3_moe_dense_linear_q4k_q8_1_x4(%token_count$41: index) launch(%token_count$42: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count$46: index, %input_size$47: index, %output_size$48: index) launch(%token_count$49: index, %input_size$50: index, %output_size$51: index, %q8_input: buffer, %weight: buffer, %output: buffer) kernel.def target(@qwen3_moe_attention_qkv_gfx11_wave32) @qwen3_moe_attention_qkv_quantized(%token_count: index) { %query_size = config.get @qwen3_moe.attention.query_size : index @@ -161,7 +64,7 @@ kernel.def target(@qwen3_moe_attention_qkv_gfx11_wave32) @qwen3_moe_attention_qk %c0 = index.constant 0 : index %valid_token = index.cmp ult, %token0, %bounded_token_count : index %safe_token = scf.select %valid_token, %token0, %c0 : index - func.apply(%value_uses_q6_index, %valid_token, %bounded_token_count, %safe_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + template.apply<@qwen3_moe.attention.qkv_quantized.body>(%value_uses_q6_index, %valid_token, %bounded_token_count, %safe_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output, %key_output, %value_output) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) kernel.return } @@ -173,7 +76,7 @@ check.case public @qwen3_moe_attention_qkv_q6_differential_case { %token_count = check.literal value(14) : index %hidden_size = check.literal value(512) : index %output_size = check.literal value(65) : index - %input_seed = check.param.seed base(0x514d4f45514b5636) count(1) : i64 + %input_seed = check.param.seed base(5858425849414112822) count(1) : i64 %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<14x512xf32> %q8_input = check.generate.fill value(0) : tensor<14x576xi8> %query_weight = check.generate.fill value(34) : tensor<65x2x144xi8> @@ -199,7 +102,7 @@ check.case public @qwen3_moe_attention_qkv_q6_differential_case { check.case public @qwen3_moe_attention_qkv_q4_differential_case { %token_count = check.literal value(14) : index %hidden_size = check.literal value(512) : index - %input_seed = check.param.seed base(0x514d4f45514b5634) count(1) : i64 + %input_seed = check.param.seed base(5858425849414112820) count(1) : i64 %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<14x512xf32> %q8_input = check.generate.fill value(0) : tensor<14x576xi8> %query_weight = check.generate.fill value(34) : tensor<65x2x144xi8> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_same_format_prefill.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_same_format_prefill.loom deleted file mode 100644 index 92e7bc805041..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/attention_qkv_same_format_prefill.loom +++ /dev/null @@ -1,191 +0,0 @@ -// Same-format Q4_K attention projections over one contiguous Q/K/V tensor. -// -// The fixed parameter layout places Q, K, and Q4_K V rows consecutively. This -// bounded candidate pairs those [5120][2048] weights with a row-interleaved -// [token][5120] raw output so one uniform 80xN grid can reuse the canonical -// dense WMMA body without buffer-valued control flow, completion counters, or -// additional workgroup storage. -amdgpu.target @qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] - -config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] - -config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -// Runs every Q4_K Q/K/V tile in one launch over the aggregate physical rows. -kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_qkv_q4_prefill_512(%token_count: index) { - %query_size = config.get @qwen3_moe.attention.query_size : index - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %combined_output_size = index.add %query_size, %key_value_output_size : index - %padded_output_size = index.add %combined_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, %combined_weight: buffer, %combined_output: buffer) { - %hidden_size = config.get @qwen3_moe.model.hidden_size : index - %query_size = config.get @qwen3_moe.attention.query_size : index - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %c2 = index.constant 2 : index - %q4 = index.constant 4 : index - %overwrite = index.constant 0 : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %combined_output_size = index.add %query_size, %key_value_output_size : index - func.apply(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -// The three segment exports form the exact aggregate-layout baseline. Their -// workgroups and memory accesses match the combined launch; only the launch -// records differ. -kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_query_q4_aggregate_prefill_512(%token_count: index) { - %query_size = config.get @qwen3_moe.attention.query_size : index - %token_capacity = config.get @qwen3_moe.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 %query_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, %combined_weight: buffer, %combined_output: buffer) { - %hidden_size = config.get @qwen3_moe.model.hidden_size : index - %query_size = config.get @qwen3_moe.attention.query_size : index - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %c2 = index.constant 2 : index - %q4 = index.constant 4 : index - %overwrite = index.constant 0 : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %combined_output_size = index.add %query_size, %key_value_output_size : index - func.apply(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_key_q4_aggregate_prefill_512(%token_count: index) { - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %token_capacity = config.get @qwen3_moe.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 %key_value_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, %combined_weight: buffer, %combined_output: buffer) { - %hidden_size = config.get @qwen3_moe.model.hidden_size : index - %query_size = config.get @qwen3_moe.attention.query_size : index - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %local_channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %c2 = index.constant 2 : index - %q4 = index.constant 4 : index - %c64 = index.constant 64 : index - %overwrite = index.constant 0 : index - %query_channel_tile_count = index.div %query_size, %c64 : index - %channel_tile = index.add %query_channel_tile_count, %local_channel_tile : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %combined_output_size = index.add %query_size, %key_value_output_size : index - func.apply(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_attention_qkv_same_format_prefill_gfx11_wave64) @qwen3_moe_attention_value_q4_aggregate_prefill_512(%token_count: index) { - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %token_capacity = config.get @qwen3_moe.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 %key_value_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, %combined_weight: buffer, %combined_output: buffer) { - %hidden_size = config.get @qwen3_moe.model.hidden_size : index - %query_size = config.get @qwen3_moe.attention.query_size : index - %key_value_size = config.get @qwen3_moe.attention.key_value_size : index - %local_channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %c2 = index.constant 2 : index - %q4 = index.constant 4 : index - %c64 = index.constant 64 : index - %overwrite = index.constant 0 : index - %query_channel_tile_count = index.div %query_size, %c64 : index - %key_value_channel_tile_count = index.div %key_value_size, %c64 : index - %key_channel_tile_end = index.add %query_channel_tile_count, %key_value_channel_tile_count : index - %channel_tile = index.add %key_channel_tile_end, %local_channel_tile : index - %key_value_output_size = index.mul %key_value_size, %c2 : index - %combined_output_size = index.add %query_size, %key_value_output_size : index - func.apply(%q4, %token_count, %hidden_size, %combined_output_size, %overwrite, %channel_tile, %token_tile, %input, %combined_weight, %combined_output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -// The production shape crosses every token and channel tile. The aggregate -// output starts with different sentinels so missing or overlapping segment -// ownership is observable even though the packed bytes are uniform. -check.case public @qwen3_moe_attention_qkv_q4_prefill_512_differential_case { - %token_count = check.literal value(512) : index - %input = check.generate.fill value(0.00390625) : tensor<512x2048xf32> - %combined_weight = check.generate.fill value(34) : tensor<5120x8x144xi8> - %expected = check.generate.fill value(-1.0) : tensor<512x5120xf32> - %actual = check.generate.fill value(1.0) : tensor<512x5120xf32> - kernel.launch @qwen3_moe_attention_query_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %expected) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - kernel.launch @qwen3_moe_attention_key_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %expected) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - kernel.launch @qwen3_moe_attention_value_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %expected) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - kernel.launch @qwen3_moe_attention_qkv_q4_prefill_512[%token_count](%token_count, %input, %combined_weight, %actual) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<512x5120xf32> - check.return -} - -check.case public @qwen3_moe_attention_qkv_q4_prefill_512_composed_benchmark_case { - %token_count = check.literal value(512) : index - %input = check.generate.fill value(0.0) : tensor<512x2048xf32> - %combined_weight = check.generate.fill value(0) : tensor<5120x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<512x5120xf32> - kernel.launch @qwen3_moe_attention_query_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - kernel.launch @qwen3_moe_attention_key_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - kernel.launch @qwen3_moe_attention_value_q4_aggregate_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - check.return -} - -check.case public @qwen3_moe_attention_qkv_q4_prefill_512_fused_benchmark_case { - %token_count = check.literal value(512) : index - %input = check.generate.fill value(0.0) : tensor<512x2048xf32> - %combined_weight = check.generate.fill value(0) : tensor<5120x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<512x5120xf32> - kernel.launch @qwen3_moe_attention_qkv_q4_prefill_512[%token_count](%token_count, %input, %combined_weight, %output) : [index](index, tensor<512x2048xf32>, tensor<5120x8x144xi8>, tensor<512x5120xf32>) - check.return -} - -check.benchmark<@qwen3_moe_attention_qkv_q4_prefill_512_differential_case> @qwen3_moe_attention_qkv_q4_prefill_512_differential - -check.benchmark<@qwen3_moe_attention_qkv_q4_prefill_512_composed_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_512_composed - -check.benchmark<@qwen3_moe_attention_qkv_q4_prefill_512_fused_benchmark_case> @qwen3_moe_attention_qkv_q4_prefill_512_fused diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_expert_dispatch.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_expert_dispatch.loom deleted file mode 100644 index b45ac4ea517d..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_expert_dispatch.loom +++ /dev/null @@ -1,498 +0,0 @@ -// Builds the GPU-resident expert work queues used by batched decode. -// -// One lane owns each configured expert, counts its assignments, and uses -// workgroup scans to assign deterministic compact offsets. Three configured -// row-count limits classify active experts into four target-selected schedule -// queues. The resulting counts are suitable for device-side indirect launch; -// route identities never cross the host boundary. -// -// assignment_ordinals is expert-contiguous. Each ordinal is the original -// flattened [token][route] position, so consumers can recover token and route -// with division and remainder by the configured route count and scatter -// results directly into the established compact routed-output layout. -// -// queue_descriptors is physically [4][expert_count][4xi32]. Each naturally -// aligned descriptor contains the expert ordinal, assignment base, row count, -// and a reserved zero field. The b128 representation keeps the ABI independent -// of model geometry while giving consumers one native-width descriptor load. -config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] - -config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -config.decl @qwen3_moe.batched_decode.schedule0_row_limit : %value: index where [range(%value, 1, 2048)] - -config.decl @qwen3_moe.batched_decode.schedule1_row_limit : %value: index where [range(%value, 1, 2048)] - -config.decl @qwen3_moe.batched_decode.schedule2_row_limit : %value: index where [range(%value, 1, 2048)] - -amdgpu.target @qwen3_moe_batched_decode_dispatch_gfx11_wave32 {subgroup_size = 32} - -// Decodes the stable batched-decode expert descriptor representation. -func.def inline @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor: vector<4xi32>) -> (index, index, index) { - %configured_expert_count = config.get @qwen3_moe.router.expert_count : index - %configured_route_count = config.get @qwen3_moe.router.route_count : index - %configured_token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %assignment_capacity = index.mul %configured_token_capacity, %configured_route_count : index - %expert_i32 = vector.extract %descriptor[0] : vector<4xi32> -> i32 - %assignment_base_i32 = vector.extract %descriptor[1] : vector<4xi32> -> i32 - %row_count_i32 = vector.extract %descriptor[2] : vector<4xi32> -> i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert, %descriptor_expert_count = index.assume %expert0, %configured_expert_count [range(%expert0, 0, 511), lt(%expert0, %configured_expert_count)] : index, index - %assignment_base0 = index.cast %assignment_base_i32 : i32 to index - %assignment_base, %descriptor_assignment_capacity = index.assume %assignment_base0, %assignment_capacity [range(%assignment_base0, 0, 65535), lt(%assignment_base0, %assignment_capacity)] : index, index - %row_count0 = index.cast %row_count_i32 : i32 to index - %row_count, %descriptor_token_capacity = index.assume %row_count0, %configured_token_capacity [range(%row_count0, 1, 2048), le(%row_count0, %configured_token_capacity)] : index, index - func.return %expert, %assignment_base, %row_count : index, index, index -} - -kernel.def target(@qwen3_moe_batched_decode_dispatch_gfx11_wave32) @qwen3_moe_build_batched_decode_expert_dispatch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { - %configured_expert_count = config.get @qwen3_moe.router.expert_count : index - %c1 = index.constant 1 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%configured_expert_count, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %assignment_ordinals: buffer, %queue_counts: buffer, %queue_descriptors: buffer) { - %configured_token_capacity0 = config.get @qwen3_moe.workload.token_capacity : index - %bounded_token_count, %configured_token_capacity = index.assume %token_count, %configured_token_capacity0 [range(%token_count, 1, 2048), le(%token_count, %configured_token_capacity0)] : index, index - %configured_route_count0 = config.get @qwen3_moe.router.route_count : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_stride, %route_row_width = index.assume %route_stride, %bounded_route_count [range(%route_stride, 1, 512), le(%bounded_route_count, %route_stride)] : index, index - %configured_expert_count0 = config.get @qwen3_moe.router.expert_count : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 32, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %schedule0_row_limit0 = config.get @qwen3_moe.batched_decode.schedule0_row_limit : index - %schedule1_row_limit0 = config.get @qwen3_moe.batched_decode.schedule1_row_limit : index - %schedule2_row_limit0 = config.get @qwen3_moe.batched_decode.schedule2_row_limit : index - %schedule0_row_limit, %schedule1_row_limit, %schedule2_row_limit = index.assume %schedule0_row_limit0, %schedule1_row_limit0, %schedule2_row_limit0 [lt(%schedule0_row_limit0, %schedule1_row_limit0), lt(%schedule1_row_limit0, %schedule2_row_limit0)] : index, index, index - %lane0 = kernel.workitem.id : index - %expert, %launch_expert_count = index.assume %lane0, %bounded_expert_count [lt(%lane0, %bounded_expert_count)] : index, index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %assignment_count = index.mul %bounded_token_count, %configured_route_count : index - %route_ids_noalias, %assignment_ordinals_noalias, %queue_counts_noalias, %queue_descriptors_noalias = buffer.assume.noalias %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors : buffer, buffer, buffer, buffer - %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> - %assignment_view = buffer.view %assignment_ordinals_noalias[%c0_offset] : buffer -> view<[%assignment_count]xi32> - %queue_count_view = buffer.view %queue_counts_noalias[%c0_offset] : buffer -> view<4xi32> - %queue_descriptor_view = buffer.view %queue_descriptors_noalias[%c0_offset] : buffer -> view<4x[%configured_expert_count]x4xi32> - %expert_i32 = index.cast %expert : index to i32 - - // Every lane walks the same tiny, cache-resident route plane. The serial - // comparisons are substantially cheaper than a host-visible sort or atomics - // and make one lane the sole owner of each expert's count. - %expert_assignment_count = scf.for %assignment = [%c0 to %assignment_count step %c1](%matched_count = %c0_i32 : i32) -> (i32) { - %token0 = index.div %assignment, %configured_route_count : index - %route0 = index.rem %assignment, %configured_route_count : index - %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index - %route_expert_i32 = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - %matches = scalar.cmpi eq, %route_expert_i32, %expert_i32 : i32 - %match_increment = scf.select %matches, %c1_i32, %c0_i32 : i32 - %next_matched_count = scalar.addi %matched_count, %match_increment : i32 - scf.yield %next_matched_count : i32 - } - - %expert_assignment_base_i32 = kernel.workgroup.scan %expert_assignment_count {direction = forward, mode = exclusive} : i32 - - // A second pass publishes stable original assignment ordinals into the - // expert-contiguous permutation. This metadata remains cache-resident for - // practical decode batches. - %published_count = scf.for %assignment = [%c0 to %assignment_count step %c1](%matched_count = %c0_i32 : i32) -> (i32) { - %token0 = index.div %assignment, %configured_route_count : index - %route0 = index.rem %assignment, %configured_route_count : index - %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index - %route_expert_i32 = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - %matches = scalar.cmpi eq, %route_expert_i32, %expert_i32 : i32 - scf.if %matches { - %compact_ordinal_i32 = scalar.addi %expert_assignment_base_i32, %matched_count : i32 - %compact_ordinal0 = index.cast %compact_ordinal_i32 : i32 to index - %bounded_compact_ordinal, %bounded_assignment_count = index.assume %compact_ordinal0, %assignment_count [range(%compact_ordinal0, 0, 65535), lt(%compact_ordinal0, %assignment_count)] : index, index - %assignment_i32 = index.cast %assignment : index to i32 - view.store %assignment_i32, %assignment_view[%bounded_compact_ordinal] : i32, view<[%assignment_count]xi32> - } - %match_increment = scf.select %matches, %c1_i32, %c0_i32 : i32 - %next_matched_count = scalar.addi %matched_count, %match_increment : i32 - scf.yield %next_matched_count : i32 - } - - // Top-k route IDs are unique within a token, so one expert owns at most the - // configured token capacity. Target-selected row limits partition that - // range without changing this producer or its ABI. - %published_count0 = index.cast %published_count : i32 to index - %published_row_count, %published_token_capacity = index.assume %published_count0, %configured_token_capacity [range(%published_count0, 0, 2048), le(%published_count0, %configured_token_capacity)] : index, index - %has_assignments = index.cmp ugt, %published_row_count, %c0 : index - %at_most_schedule0 = index.cmp ule, %published_row_count, %schedule0_row_limit : index - %above_schedule0 = index.cmp ugt, %published_row_count, %schedule0_row_limit : index - %at_most_schedule1 = index.cmp ule, %published_row_count, %schedule1_row_limit : index - %above_schedule1 = index.cmp ugt, %published_row_count, %schedule1_row_limit : index - %at_most_schedule2 = index.cmp ule, %published_row_count, %schedule2_row_limit : index - %is_schedule0 = scalar.andi %has_assignments, %at_most_schedule0 : i1 - %is_schedule1 = scalar.andi %above_schedule0, %at_most_schedule1 : i1 - %is_schedule2 = scalar.andi %above_schedule1, %at_most_schedule2 : i1 - %is_schedule3 = index.cmp ugt, %published_row_count, %schedule2_row_limit : index - %schedule0_i32 = scf.select %is_schedule0, %c1_i32, %c0_i32 : i32 - %schedule1_i32 = scf.select %is_schedule1, %c1_i32, %c0_i32 : i32 - %schedule2_i32 = scf.select %is_schedule2, %c1_i32, %c0_i32 : i32 - %schedule3_i32 = scf.select %is_schedule3, %c1_i32, %c0_i32 : i32 - - %schedule0_ordinal_i32 = kernel.workgroup.scan %schedule0_i32 {direction = forward, mode = exclusive} : i32 - %schedule1_ordinal_i32 = kernel.workgroup.scan %schedule1_i32 {direction = forward, mode = exclusive} : i32 - %schedule2_ordinal_i32 = kernel.workgroup.scan %schedule2_i32 {direction = forward, mode = exclusive} : i32 - %schedule3_ordinal_i32 = kernel.workgroup.scan %schedule3_i32 {direction = forward, mode = exclusive} : i32 - %schedule0_count = kernel.workgroup.reduce %schedule0_i32 : i32 - %schedule1_count = kernel.workgroup.reduce %schedule1_i32 : i32 - %schedule2_count = kernel.workgroup.reduce %schedule2_i32 : i32 - %schedule3_count = kernel.workgroup.reduce %schedule3_i32 : i32 - %descriptor = vector.from_elements %expert_i32, %expert_assignment_base_i32, %published_count, %c0_i32 : vector<4xi32> - - // Inactive lanes use ordinal zero so every physical origin is in bounds; - // the vector masks suppress their stores. This keeps the descriptor write a - // native b128 operation without placing it inside divergent control flow. - %safe_schedule0_ordinal_i32 = scf.select %is_schedule0, %schedule0_ordinal_i32, %c0_i32 : i32 - %schedule0_ordinal0 = index.cast %safe_schedule0_ordinal_i32 : i32 to index - %schedule0_ordinal, %schedule0_queue_capacity = index.assume %schedule0_ordinal0, %configured_expert_count [range(%schedule0_ordinal0, 0, 511), lt(%schedule0_ordinal0, %configured_expert_count)] : index, index - %schedule0_mask = vector.splat %is_schedule0 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[0, %schedule0_ordinal, 0], %schedule0_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - %safe_schedule1_ordinal_i32 = scf.select %is_schedule1, %schedule1_ordinal_i32, %c0_i32 : i32 - %schedule1_ordinal0 = index.cast %safe_schedule1_ordinal_i32 : i32 to index - %schedule1_ordinal, %schedule1_queue_capacity = index.assume %schedule1_ordinal0, %configured_expert_count [range(%schedule1_ordinal0, 0, 511), lt(%schedule1_ordinal0, %configured_expert_count)] : index, index - %schedule1_mask = vector.splat %is_schedule1 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[1, %schedule1_ordinal, 0], %schedule1_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - %safe_schedule2_ordinal_i32 = scf.select %is_schedule2, %schedule2_ordinal_i32, %c0_i32 : i32 - %schedule2_ordinal0 = index.cast %safe_schedule2_ordinal_i32 : i32 to index - %schedule2_ordinal, %schedule2_queue_capacity = index.assume %schedule2_ordinal0, %configured_expert_count [range(%schedule2_ordinal0, 0, 511), lt(%schedule2_ordinal0, %configured_expert_count)] : index, index - %schedule2_mask = vector.splat %is_schedule2 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[2, %schedule2_ordinal, 0], %schedule2_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - %safe_schedule3_ordinal_i32 = scf.select %is_schedule3, %schedule3_ordinal_i32, %c0_i32 : i32 - %schedule3_ordinal0 = index.cast %safe_schedule3_ordinal_i32 : i32 to index - %schedule3_ordinal, %schedule3_queue_capacity = index.assume %schedule3_ordinal0, %configured_expert_count [range(%schedule3_ordinal0, 0, 511), lt(%schedule3_ordinal0, %configured_expert_count)] : index, index - %schedule3_mask = vector.splat %is_schedule3 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[3, %schedule3_ordinal, 0], %schedule3_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - - %is_lane_zero = index.cmp eq, %expert, %c0 : index - scf.if %is_lane_zero { - view.store %schedule0_count, %queue_count_view[0] : i32, view<4xi32> - view.store %schedule1_count, %queue_count_view[1] : i32, view<4xi32> - view.store %schedule2_count, %queue_count_view[2] : i32, view<4xi32> - view.store %schedule3_count, %queue_count_view[3] : i32, view<4xi32> - } - kernel.return -} - -// Serial specification oracle for exact differential tests. Its ownership and -// control flow deliberately differ from the parallel scan implementation: -// one workitem visits experts in order and carries all compact offsets and -// queue tails explicitly. -kernel.def target(@qwen3_moe_batched_decode_dispatch_gfx11_wave32) @qwen3_moe_build_batched_decode_expert_dispatch_reference(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { - %c1 = index.constant 1 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c1, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %assignment_ordinals: buffer, %queue_counts: buffer, %queue_descriptors: buffer) { - %configured_token_capacity0 = config.get @qwen3_moe.workload.token_capacity : index - %bounded_token_count, %configured_token_capacity = index.assume %token_count, %configured_token_capacity0 [range(%token_count, 1, 2048), le(%token_count, %configured_token_capacity0)] : index, index - %configured_route_count0 = config.get @qwen3_moe.router.route_count : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_stride, %route_row_width = index.assume %route_stride, %bounded_route_count [range(%route_stride, 1, 512), le(%bounded_route_count, %route_stride)] : index, index - %configured_expert_count0 = config.get @qwen3_moe.router.expert_count : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 32, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %schedule0_row_limit0 = config.get @qwen3_moe.batched_decode.schedule0_row_limit : index - %schedule1_row_limit0 = config.get @qwen3_moe.batched_decode.schedule1_row_limit : index - %schedule2_row_limit0 = config.get @qwen3_moe.batched_decode.schedule2_row_limit : index - %schedule0_row_limit, %schedule1_row_limit, %schedule2_row_limit = index.assume %schedule0_row_limit0, %schedule1_row_limit0, %schedule2_row_limit0 [lt(%schedule0_row_limit0, %schedule1_row_limit0), lt(%schedule1_row_limit0, %schedule2_row_limit0)] : index, index, index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %assignment_count = index.mul %bounded_token_count, %configured_route_count : index - %route_ids_noalias, %assignment_ordinals_noalias, %queue_counts_noalias, %queue_descriptors_noalias = buffer.assume.noalias %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors : buffer, buffer, buffer, buffer - %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> - %assignment_view = buffer.view %assignment_ordinals_noalias[%c0_offset] : buffer -> view<[%assignment_count]xi32> - %queue_count_view = buffer.view %queue_counts_noalias[%c0_offset] : buffer -> view<4xi32> - %queue_descriptor_view = buffer.view %queue_descriptors_noalias[%c0_offset] : buffer -> view<4x[%configured_expert_count]x4xi32> - - %final_assignment_base, %final_schedule0_count, %final_schedule1_count, %final_schedule2_count, %final_schedule3_count = scf.for %expert = [%c0 to %bounded_expert_count step %c1](%assignment_base = %c0_i32 : i32, %schedule0_count = %c0_i32 : i32, %schedule1_count = %c0_i32 : i32, %schedule2_count = %c0_i32 : i32, %schedule3_count = %c0_i32 : i32) -> (i32, i32, i32, i32, i32) { - %expert_i32 = index.cast %expert : index to i32 - %expert_assignment_count = scf.for %assignment = [%c0 to %assignment_count step %c1](%matched_count = %c0_i32 : i32) -> (i32) { - %token0 = index.div %assignment, %configured_route_count : index - %route0 = index.rem %assignment, %configured_route_count : index - %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index - %route_expert_i32 = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - %matches = scalar.cmpi eq, %route_expert_i32, %expert_i32 : i32 - scf.if %matches { - %compact_ordinal_i32 = scalar.addi %assignment_base, %matched_count : i32 - %compact_ordinal0 = index.cast %compact_ordinal_i32 : i32 to index - %bounded_compact_ordinal, %bounded_assignment_count = index.assume %compact_ordinal0, %assignment_count [range(%compact_ordinal0, 0, 65535), lt(%compact_ordinal0, %assignment_count)] : index, index - %assignment_i32 = index.cast %assignment : index to i32 - view.store %assignment_i32, %assignment_view[%bounded_compact_ordinal] : i32, view<[%assignment_count]xi32> - } - %match_increment = scf.select %matches, %c1_i32, %c0_i32 : i32 - %next_matched_count = scalar.addi %matched_count, %match_increment : i32 - scf.yield %next_matched_count : i32 - } - - %expert_assignment_count0 = index.cast %expert_assignment_count : i32 to index - %expert_row_count, %expert_token_capacity = index.assume %expert_assignment_count0, %configured_token_capacity [range(%expert_assignment_count0, 0, 2048), le(%expert_assignment_count0, %configured_token_capacity)] : index, index - %has_assignments = index.cmp ugt, %expert_row_count, %c0 : index - %at_most_schedule0 = index.cmp ule, %expert_row_count, %schedule0_row_limit : index - %above_schedule0 = index.cmp ugt, %expert_row_count, %schedule0_row_limit : index - %at_most_schedule1 = index.cmp ule, %expert_row_count, %schedule1_row_limit : index - %above_schedule1 = index.cmp ugt, %expert_row_count, %schedule1_row_limit : index - %at_most_schedule2 = index.cmp ule, %expert_row_count, %schedule2_row_limit : index - %is_schedule0 = scalar.andi %has_assignments, %at_most_schedule0 : i1 - %is_schedule1 = scalar.andi %above_schedule0, %at_most_schedule1 : i1 - %is_schedule2 = scalar.andi %above_schedule1, %at_most_schedule2 : i1 - %is_schedule3 = index.cmp ugt, %expert_row_count, %schedule2_row_limit : index - %descriptor = vector.from_elements %expert_i32, %assignment_base, %expert_assignment_count, %c0_i32 : vector<4xi32> - - // The scalar oracle retains its loop-carried queue tails. Masked b128 - // stores express conditional publication without divergent branch regions, - // while safe ordinal zero keeps inactive physical origins in bounds. - %safe_schedule0_count = scf.select %is_schedule0, %schedule0_count, %c0_i32 : i32 - %schedule0_ordinal0 = index.cast %safe_schedule0_count : i32 to index - %schedule0_ordinal, %schedule0_capacity = index.assume %schedule0_ordinal0, %configured_expert_count [range(%schedule0_ordinal0, 0, 511), lt(%schedule0_ordinal0, %configured_expert_count)] : index, index - %schedule0_mask = vector.splat %is_schedule0 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[0, %schedule0_ordinal, 0], %schedule0_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - %safe_schedule1_count = scf.select %is_schedule1, %schedule1_count, %c0_i32 : i32 - %schedule1_ordinal0 = index.cast %safe_schedule1_count : i32 to index - %schedule1_ordinal, %schedule1_capacity = index.assume %schedule1_ordinal0, %configured_expert_count [range(%schedule1_ordinal0, 0, 511), lt(%schedule1_ordinal0, %configured_expert_count)] : index, index - %schedule1_mask = vector.splat %is_schedule1 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[1, %schedule1_ordinal, 0], %schedule1_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - %safe_schedule2_count = scf.select %is_schedule2, %schedule2_count, %c0_i32 : i32 - %schedule2_ordinal0 = index.cast %safe_schedule2_count : i32 to index - %schedule2_ordinal, %schedule2_capacity = index.assume %schedule2_ordinal0, %configured_expert_count [range(%schedule2_ordinal0, 0, 511), lt(%schedule2_ordinal0, %configured_expert_count)] : index, index - %schedule2_mask = vector.splat %is_schedule2 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[2, %schedule2_ordinal, 0], %schedule2_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - %safe_schedule3_count = scf.select %is_schedule3, %schedule3_count, %c0_i32 : i32 - %schedule3_ordinal0 = index.cast %safe_schedule3_count : i32 to index - %schedule3_ordinal, %schedule3_capacity = index.assume %schedule3_ordinal0, %configured_expert_count [range(%schedule3_ordinal0, 0, 511), lt(%schedule3_ordinal0, %configured_expert_count)] : index, index - %schedule3_mask = vector.splat %is_schedule3 : vector<4xi1> - vector.store.mask %descriptor, %queue_descriptor_view[3, %schedule3_ordinal, 0], %schedule3_mask : vector<4xi32>, view<4x[%configured_expert_count]x4xi32>, vector<4xi1> - - %schedule0_increment = scf.select %is_schedule0, %c1_i32, %c0_i32 : i32 - %schedule1_increment = scf.select %is_schedule1, %c1_i32, %c0_i32 : i32 - %schedule2_increment = scf.select %is_schedule2, %c1_i32, %c0_i32 : i32 - %schedule3_increment = scf.select %is_schedule3, %c1_i32, %c0_i32 : i32 - %next_assignment_base = scalar.addi %assignment_base, %expert_assignment_count : i32 - %next_schedule0_count = scalar.addi %schedule0_count, %schedule0_increment : i32 - %next_schedule1_count = scalar.addi %schedule1_count, %schedule1_increment : i32 - %next_schedule2_count = scalar.addi %schedule2_count, %schedule2_increment : i32 - %next_schedule3_count = scalar.addi %schedule3_count, %schedule3_increment : i32 - scf.yield %next_assignment_base, %next_schedule0_count, %next_schedule1_count, %next_schedule2_count, %next_schedule3_count : i32, i32, i32, i32, i32 - } - - view.store %final_schedule0_count, %queue_count_view[0] : i32, view<4xi32> - view.store %final_schedule1_count, %queue_count_view[1] : i32, view<4xi32> - view.store %final_schedule2_count, %queue_count_view[2] : i32, view<4xi32> - view.store %final_schedule3_count, %queue_count_view[3] : i32, view<4xi32> - kernel.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_singleton_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_pair_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(64) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_three_row_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(43) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_four_row_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(32) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_five_to_six_row_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(25) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_fourteen_to_fifteen_row_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(9) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_coherent_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<16x8xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<128xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<128xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<128xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x128x4xi32> - check.return -} - -// A second model geometry proves that expert, route, batch, descriptor, and -// schedule boundaries are configuration rather than Qwen3-30B constants. -check.case public @qwen3_moe_batched_decode_expert_dispatch_configurable_geometry_case { - %token_count = check.literal value(8) : index - %route_count = check.literal value(4) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(32) : index - %route_ids = check.generate.iota offset(0) step(1) period(13) : tensor<8x4xi32> - %actual_assignments = check.generate.fill value(-1) : tensor<32xi32> - %actual_counts = check.generate.fill value(-1) : tensor<4xi32> - %actual_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> - %expected_assignments = check.generate.fill value(-1) : tensor<32xi32> - %expected_counts = check.generate.fill value(-1) : tensor<4xi32> - %expected_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_assignments, %actual_counts, %actual_descriptors) : [index, index, index, index](index, index, index, index, tensor<8x4xi32>, tensor<32xi32>, tensor<4xi32>, tensor<4x32x4xi32>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch_reference[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_assignments, %expected_counts, %expected_descriptors) : [index, index, index, index](index, index, index, index, tensor<8x4xi32>, tensor<32xi32>, tensor<4xi32>, tensor<4x32x4xi32>) - check.expect.equal actual(%actual_assignments) expected(%expected_assignments) : tensor<32xi32> - check.expect.equal actual(%actual_counts) expected(%expected_counts) : tensor<4xi32> - check.expect.equal actual(%actual_descriptors) expected(%expected_descriptors) : tensor<4x32x4xi32> - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_diverse_benchmark_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<16x8xi32> - %assignment_ordinals = check.generate.fill value(-1) : tensor<128xi32> - %queue_counts = check.generate.fill value(-1) : tensor<4xi32> - %queue_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_coherent_benchmark_case { - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<16x8xi32> - %assignment_ordinals = check.generate.fill value(-1) : tensor<128xi32> - %queue_counts = check.generate.fill value(-1) : tensor<4xi32> - %queue_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - check.return -} - -check.case public @qwen3_moe_batched_decode_expert_dispatch_configurable_benchmark_case { - %token_count = check.literal value(8) : index - %route_count = check.literal value(4) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(32) : index - %route_ids = check.generate.iota offset(0) step(1) period(13) : tensor<8x4xi32> - %assignment_ordinals = check.generate.fill value(-1) : tensor<32xi32> - %queue_counts = check.generate.fill value(-1) : tensor<4xi32> - %queue_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<8x4xi32>, tensor<32xi32>, tensor<4xi32>, tensor<4x32x4xi32>) - check.return -} - -check.benchmark<@qwen3_moe_batched_decode_expert_dispatch_diverse_benchmark_case> @qwen3_moe_batched_decode_expert_dispatch_diverse - -check.benchmark<@qwen3_moe_batched_decode_expert_dispatch_coherent_benchmark_case> @qwen3_moe_batched_decode_expert_dispatch_coherent - -check.benchmark<@qwen3_moe_batched_decode_expert_dispatch_configurable_benchmark_case> @qwen3_moe_batched_decode_expert_dispatch_configurable diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_gate_up_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_gate_up_q4k.loom deleted file mode 100644 index 891d32890b12..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/batched_decode_gate_up_q4k.loom +++ /dev/null @@ -1,315 +0,0 @@ -// Expert-grouped batched-decode gate/up consumers for raw GGUF Q4_K weights. -// -// The assignment producer owns routing and publishes expert-contiguous -// assignment ordinals plus aligned descriptors. This file owns row-count -// schedules that consume those descriptors without inspecting route IDs. The -// first provider is exact M_e=2: each wave reuses one expert/channel weight -// row across two independently quantized activation rows and scatters the -// fused SwiGLU results back to their original [token][route] assignments. -config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] - -config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] - -config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] - -config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -// Static fallback capacity for the exact-two-row queue. Device-side indirect -// launch replaces this upper bound with the producer's exact queue count. -config.decl @qwen3_moe.batched_decode.rows2_descriptor_capacity : %value: index where [range(%value, 1, 512)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) - -func.decl @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) - -func.decl @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor: vector<4xi32>) -> (index, index, index) - -kernel.decl @qwen3_moe_build_batched_decode_expert_dispatch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %assignment_ordinals: buffer, %queue_counts: buffer, %queue_descriptors: buffer) - -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -// Contracts one expert/channel Q4_K row against two independently quantized -// activation rows. Splitting gate and up into separate contractions keeps each -// recurrence small while preserving the important invariant: every weight -// byte loaded by this helper contributes to both rows. -func.def inline @qwen3_moe_batched_decode_q4k_rows2_lane(%input_size: index, %weight: buffer, %q8_input: buffer, %weight_row_byte_base: offset, %q8_row0_byte_base: offset, %q8_row1_byte_base: offset, %lane: index) -> (f32, f32) { - %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 - %c256 = index.constant 256 : index - %c1023 = index.constant 1023 : index - %c1024 = index.constant 1024 : index - %q4_block_bytes = index.constant 144 : offset - %q8_group_bytes = index.constant 144 : offset - %q8_payload_byte_add = index.constant 16 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %input_size, %c256 : index - %q8_group_count = index.div %input_size, %c128 : index - %padded_input_size = index.add %input_size, %c1023 : index - %iteration_count = index.div %padded_input_size, %c1024 : index - %q4_group_pair0 = index.div %lane, %c2 : index - %q4_group_pair1 = index.rem %q4_group_pair0, %c4 : index - %q4_group_pair = index.assume %q4_group_pair1 [range(%q4_group_pair1, 0, 3)] : index - %q4_half0 = index.rem %lane, %c2 : index - %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index - %lane_q4_block = index.div %lane, %c8 : index - %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index - %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index - %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index - %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index - %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index - %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index - %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index - %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index - %q8_half_word_add = index.mul %q4_half, %c4 : index - %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index - %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index - %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index - %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index - %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index - %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index - %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index - %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index - - %row0_acc, %row1_acc = scf.for %iteration = [%c0 to %iteration_count step %c1](%row0_iter = %c0_f32 : f32, %row1_iter = %c0_f32 : f32) -> (f32, f32) unroll { - %iteration_q4_block = index.mul %iteration, %c4 : index - %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index - %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index - %row0_contribution, %row1_contribution = scf.if %valid_q4_block -> (f32, f32) { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %q8_block_group_base = index.mul %q4_block, %c2 : index - %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index - %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index - %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset - %q8_group0_byte_base = index.add %q8_row0_byte_base, %q8_group_byte_add : offset - %q8_group1_byte_base = index.add %q8_row1_byte_base, %q8_group_byte_add : offset - %q8_payload0_byte_base = index.add %q8_group0_byte_base, %q8_payload_byte_add : offset - %q8_payload1_byte_base = index.add %q8_group1_byte_base, %q8_payload_byte_add : offset - %q8_ds0_view = buffer.view %q8_input[%q8_group0_byte_base] : buffer -> view<8xf16> - %q8_ds1_view = buffer.view %q8_input[%q8_group1_byte_base] : buffer -> view<8xf16> - %q8_words0_view = buffer.view %q8_input[%q8_payload0_byte_base] : buffer -> view<32xi32> - %q8_words1_view = buffer.view %q8_input[%q8_payload1_byte_base] : buffer -> view<32xi32> - %q8_ds0 = vector.load %q8_ds0_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> - %q8_ds1 = vector.load %q8_ds1_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> - %q8_low_d0_f16 = vector.extract %q8_ds0[0] : vector<4xf16> -> f16 - %q8_low_s0_f16 = vector.extract %q8_ds0[1] : vector<4xf16> -> f16 - %q8_high_d0_f16 = vector.extract %q8_ds0[2] : vector<4xf16> -> f16 - %q8_high_s0_f16 = vector.extract %q8_ds0[3] : vector<4xf16> -> f16 - %q8_low_d1_f16 = vector.extract %q8_ds1[0] : vector<4xf16> -> f16 - %q8_low_s1_f16 = vector.extract %q8_ds1[1] : vector<4xf16> -> f16 - %q8_high_d1_f16 = vector.extract %q8_ds1[2] : vector<4xf16> -> f16 - %q8_high_s1_f16 = vector.extract %q8_ds1[3] : vector<4xf16> -> f16 - %q8_low_d0 = scalar.extf %q8_low_d0_f16 : f16 to f32 - %q8_low_s0 = scalar.extf %q8_low_s0_f16 : f16 to f32 - %q8_high_d0 = scalar.extf %q8_high_d0_f16 : f16 to f32 - %q8_high_s0 = scalar.extf %q8_high_s0_f16 : f16 to f32 - %q8_low_d1 = scalar.extf %q8_low_d1_f16 : f16 to f32 - %q8_low_s1 = scalar.extf %q8_low_s1_f16 : f16 to f32 - %q8_high_d1 = scalar.extf %q8_high_d1_f16 : f16 to f32 - %q8_high_s1 = scalar.extf %q8_high_s1_f16 : f16 to f32 - %q8_low_words0 = vector.load %q8_words0_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> - %q8_high_words0 = vector.load %q8_words0_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> - %q8_low_words1 = vector.load %q8_words1_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> - %q8_high_words1 = vector.load %q8_words1_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> - %q8_low_values0 = vector.bitcast %q8_low_words0 : vector<4xi32> to vector<16xi8> - %q8_high_values0 = vector.bitcast %q8_high_words0 : vector<4xi32> to vector<16xi8> - %q8_low_values1 = vector.bitcast %q8_low_words1 : vector<4xi32> to vector<16xi8> - %q8_high_values1 = vector.bitcast %q8_high_words1 : vector<4xi32> to vector<16xi8> - %q4_block_byte_add = index.scale %q4_block, %q4_block_bytes : index, offset -> offset - %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset - %header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> - %header_words = vector.load %header_view[0] : view<4xi32> -> vector<4xi32> - %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %q4_block, %q4_group_pair, %q4_half, %header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) - %row0_low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values0, %q8_low_d0, %q8_low_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %row0_high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values0, %q8_high_d0, %q8_high_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %row1_low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values1, %q8_low_d1, %q8_low_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %row1_high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values1, %q8_high_d1, %q8_high_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %row0_pair = scalar.addf %row0_low, %row0_high : f32 - %row1_pair = scalar.addf %row1_low, %row1_high : f32 - scf.yield %row0_pair, %row1_pair : f32, f32 - } else { - scf.yield %c0_f32, %c0_f32 : f32, f32 - } - %row0_next = scalar.addf %row0_iter, %row0_contribution : f32 - %row1_next = scalar.addf %row1_iter, %row1_contribution : f32 - scf.yield %row0_next, %row1_next : f32, f32 - } - func.return %row0_acc, %row1_acc : f32, f32 -} - -amdgpu.target @qwen3_moe_batched_decode_gate_up_gfx11_wave32 {subgroup_size = 32} - -// Processes the exact-two-row descriptor queue. The static launch covers a -// config-derived safe descriptor capacity and guards inactive workgroups. A -// composed device launcher instead supplies the exact descriptor count as -// both launch geometry and an ABI fact. -kernel.def target(@qwen3_moe_batched_decode_gate_up_gfx11_wave32) @qwen3_moe_batched_decode_gate_up_q4k_rows2(%descriptor_count: index, %queue_ordinal: index, %token_count: index, %route_count: index, %expert_count: index, %output_size: index) { - %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %descriptor_capacity = config.get @qwen3_moe.batched_decode.rows2_descriptor_capacity : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %workgroup_size = index.constant 128 : index - %padded_output_size = index.add %configured_output_size, %c3 : index - %channel_workgroup_count = index.div %padded_output_size, %c4 : index - kernel.launch.config workgroups(%channel_workgroup_count, %descriptor_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%descriptor_count: index, %queue_ordinal: index, %token_count: index, %route_count: index, %expert_count: index, %output_size: index, %queue_descriptors: buffer, %assignment_ordinals: buffer, %q8_input: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { - %configured_token_capacity0 = config.get @qwen3_moe.workload.token_capacity : index - %bounded_token_count, %configured_token_capacity = index.assume %token_count, %configured_token_capacity0 [range(%token_count, 1, 2048), eq(%token_count, %configured_token_capacity0)] : index, index - %configured_route_count0 = config.get @qwen3_moe.router.route_count : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 32), eq(%route_count, %configured_route_count0)] : index, index - %configured_expert_count0 = config.get @qwen3_moe.router.expert_count : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 32, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %configured_weight_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_weight_expert_count, %model_expert_count = index.assume %configured_weight_expert_count0, %configured_expert_count [eq(%configured_weight_expert_count0, %configured_expert_count)] : index, index - %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %bounded_queue_ordinal = index.assume %queue_ordinal [range(%queue_ordinal, 0, 3)] : index - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %descriptor_ordinal0 = kernel.workgroup.id : index - %channel_workgroup = kernel.workgroup.id : index - %subgroup = kernel.subgroup.id : index - %lane0 = kernel.subgroup.lane.id : index - %lane = index.assume %lane0 [range(%lane0, 0, 31)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c1_byte = index.constant 1 : offset - %q4_block_bytes = index.constant 144 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_offset = index.constant 0 : offset - %assignment_count = index.mul %bounded_token_count, %configured_route_count : index - %assignment_capacity = index.mul %configured_token_capacity, %configured_route_count : index - %assignment_limited_descriptor_capacity0 = index.div %assignment_capacity, %c2 : index - %descriptor_capacity0 = index.min %configured_expert_count, %assignment_limited_descriptor_capacity0 : index - %configured_descriptor_capacity0 = config.get @qwen3_moe.batched_decode.rows2_descriptor_capacity : index - %descriptor_capacity, %configured_descriptor_capacity = index.assume %descriptor_capacity0, %configured_descriptor_capacity0 [range(%descriptor_capacity0, 1, 512), eq(%descriptor_capacity0, %configured_descriptor_capacity0)] : index, index - %bounded_descriptor_count, %queue_descriptor_capacity = index.assume %descriptor_count, %configured_descriptor_capacity [range(%descriptor_count, 1, 512), le(%descriptor_count, %configured_descriptor_capacity)] : index, index - %valid_descriptor = index.cmp ult, %descriptor_ordinal0, %bounded_descriptor_count : index - %safe_descriptor_ordinal0 = scf.select %valid_descriptor, %descriptor_ordinal0, %c0 : index - %descriptor_ordinal, %launch_descriptor_count = index.assume %safe_descriptor_ordinal0, %bounded_descriptor_count [lt(%safe_descriptor_ordinal0, %bounded_descriptor_count)] : index, index - %q4_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.mul %q4_block_count, %q4_block_bytes : index - %weight_expert_bytes = index.mul %bounded_output_size, %weight_row_bytes : index - %q8_group_count = index.div %input_size, %c128 : index - %q8_bytes_per_token = index.mul %q8_group_count, %q4_block_bytes : index - %queue_descriptors_noalias, %assignment_ordinals_noalias, %q8_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %queue_descriptors, %assignment_ordinals, %q8_input, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer - %queue_descriptor_view = buffer.view %queue_descriptors_noalias[%c0_offset] : buffer -> view<4x[%configured_expert_count]x4xi32> - %assignment_view = buffer.view %assignment_ordinals_noalias[%c0_offset] : buffer -> view<[%assignment_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf32> - %descriptor = vector.load %queue_descriptor_view[%bounded_queue_ordinal, %descriptor_ordinal, 0] : view<4x[%configured_expert_count]x4xi32> -> vector<4xi32> - %expert, %assignment_base0, %row_count0 = func.call @qwen3_moe_unpack_batched_decode_expert_descriptor(%descriptor) : (vector<4xi32>) -> (index, index, index) - %assignment_base, %consumer_assignment_count = index.assume %assignment_base0, %assignment_count [lt(%assignment_base0, %assignment_count)] : index, index - %row_count = index.assume %row_count0 [range(%row_count0, 2, 2)] : index - %assignment1_ordinal0 = index.add %assignment_base, %c1 : index - %assignment1_ordinal, %pair_assignment_count = index.assume %assignment1_ordinal0, %assignment_count [lt(%assignment1_ordinal0, %assignment_count)] : index, index - %assignment0_i32 = view.load %assignment_view[%assignment_base] : view<[%assignment_count]xi32> -> i32 - %assignment1_i32 = view.load %assignment_view[%assignment1_ordinal] : view<[%assignment_count]xi32> -> i32 - %assignment0_index0 = index.cast %assignment0_i32 : i32 to index - %assignment1_index0 = index.cast %assignment1_i32 : i32 to index - %assignment0, %assignment1, %routed_assignment_count = index.assume %assignment0_index0, %assignment1_index0, %assignment_count [range(%assignment0_index0, 0, 65535), range(%assignment1_index0, 0, 65535), lt(%assignment0_index0, %assignment_count), lt(%assignment1_index0, %assignment_count)] : index, index, index - %token0_index0 = index.div %assignment0, %configured_route_count : index - %token1_index0 = index.div %assignment1, %configured_route_count : index - %token0, %token1, %q8_token_count = index.assume %token0_index0, %token1_index0, %bounded_token_count [lt(%token0_index0, %bounded_token_count), lt(%token1_index0, %bounded_token_count)] : index, index, index - %channel_base = index.mul %channel_workgroup, %c4 : index - %channel0 = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel0, %bounded_output_size : index - %safe_channel0 = scf.select %valid_channel, %channel0, %c0 : index - %channel, %weight_output_size = index.assume %safe_channel0, %bounded_output_size [lt(%safe_channel0, %bounded_output_size)] : index, index - %expert_byte_base = index.mul %expert, %weight_expert_bytes : index - %channel_byte_add = index.mul %channel, %weight_row_bytes : index - %row_byte_index = index.add %expert_byte_base, %channel_byte_add : index - %row_byte_base = index.scale %row_byte_index, %c1_byte : index, offset -> offset - %q8_token0_byte_index = index.mul %token0, %q8_bytes_per_token : index - %q8_token1_byte_index = index.mul %token1, %q8_bytes_per_token : index - %q8_token0_byte_base = index.scale %q8_token0_byte_index, %c1_byte : index, offset -> offset - %q8_token1_byte_base = index.scale %q8_token1_byte_index, %c1_byte : index, offset -> offset - %gate0_acc, %gate1_acc = func.call @qwen3_moe_batched_decode_q4k_rows2_lane(%input_size, %gate_noalias, %q8_noalias, %row_byte_base, %q8_token0_byte_base, %q8_token1_byte_base, %lane) : (index, buffer, buffer, offset, offset, offset, index) -> (f32, f32) - %up0_acc, %up1_acc = func.call @qwen3_moe_batched_decode_q4k_rows2_lane(%input_size, %up_noalias, %q8_noalias, %row_byte_base, %q8_token0_byte_base, %q8_token1_byte_base, %lane) : (index, buffer, buffer, offset, offset, offset, index) -> (f32, f32) - %gate0_dot = kernel.subgroup.reduce %gate0_acc : f32 - %up0_dot = kernel.subgroup.reduce %up0_acc : f32 - %gate1_dot = kernel.subgroup.reduce %gate1_acc : f32 - %up1_dot = kernel.subgroup.reduce %up1_acc : f32 - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %valid_channel_descriptor = scalar.andi %valid_channel, %valid_descriptor : i1 - %writes_output = scalar.andi %valid_channel_descriptor, %is_lane_zero : i1 - scf.if %writes_output { - %gate0_silu = scalar.siluf %gate0_dot : f32 - %gate1_silu = scalar.siluf %gate1_dot : f32 - %result0 = scalar.mulf %gate0_silu, %up0_dot : f32 - %result1 = scalar.mulf %gate1_silu, %up1_dot : f32 - view.store %result0, %output_view[%assignment0, %channel] : f32, view<[%assignment_count]x[%bounded_output_size]xf32> - view.store %result1, %output_view[%assignment1, %channel] : f32, view<[%assignment_count]x[%bounded_output_size]xf32> - } - kernel.return -} - -// Two identical top-8 route rows produce eight exact-two-row descriptors. -// The direct route-centric provider remains the numerical oracle. -check.case public @qwen3_moe_batched_decode_gate_up_q4k_rows2_differential_case { - %descriptor_count = check.literal value(8) : index - %queue_ordinal = check.literal value(1) : index - %token_count = check.literal value(2) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(32) : index - %output_size = check.literal value(32) : index - %input = check.generate.iota offset(-0.5) step(0.001953125) period(1024) : tensor<2x512xf32> - %q8_input = check.generate.fill value(0) : tensor<2x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<2x8xi32> - %assignment_ordinals = check.generate.fill value(-1) : tensor<16xi32> - %queue_counts = check.generate.fill value(-1) : tensor<4xi32> - %queue_descriptors = check.generate.fill value(-1) : tensor<4x32x4xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<32x32x2x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<32x32x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<2x8x32xf32> - %actual = check.generate.fill value(1.0) : tensor<2x8x32xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x512xf32>, tensor<2x576xi8>) - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<2x8xi32>, tensor<16xi32>, tensor<4xi32>, tensor<4x32x4xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<2x576xi8>, tensor<2x8xi32>, tensor<32x32x2x144xi8>, tensor<32x32x2x144xi8>, tensor<2x8x32xf32>) - kernel.launch @qwen3_moe_batched_decode_gate_up_q4k_rows2[%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size](%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size, %queue_descriptors, %assignment_ordinals, %q8_input, %gate_weight, %up_weight, %actual) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<4x32x4xi32>, tensor<16xi32>, tensor<2x576xi8>, tensor<32x32x2x144xi8>, tensor<32x32x2x144xi8>, tensor<2x8x32xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x8x32xf32> - check.return -} - -// Production dimensions with a uniform two-row-per-expert distribution. -// Timing includes the assignment producer so the row schedule cannot hide its -// routing preparation cost. -check.case public @qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark_case { - %descriptor_count = check.literal value(64) : index - %queue_ordinal = check.literal value(1) : index - %token_count = check.literal value(16) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<16x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(64) : tensor<16x8xi32> - %assignment_ordinals = check.generate.fill value(-1) : tensor<128xi32> - %queue_counts = check.generate.fill value(-1) : tensor<4xi32> - %queue_descriptors = check.generate.fill value(-1) : tensor<4x128x4xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<16x8x768xf32> - kernel.launch @qwen3_moe_build_batched_decode_expert_dispatch[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %assignment_ordinals, %queue_counts, %queue_descriptors) : [index, index, index, index](index, index, index, index, tensor<16x8xi32>, tensor<128xi32>, tensor<4xi32>, tensor<4x128x4xi32>) - kernel.launch @qwen3_moe_batched_decode_gate_up_q4k_rows2[%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size](%descriptor_count, %queue_ordinal, %token_count, %route_count, %expert_count, %output_size, %queue_descriptors, %assignment_ordinals, %q8_input, %gate_weight, %up_weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<4x128x4xi32>, tensor<128xi32>, tensor<16x2304xi8>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<16x8x768xf32>) - check.return -} - -check.benchmark<@qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark_case> @qwen3_moe_batched_decode_gate_up_q4k_rows2_benchmark diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom index 95fe546dcdbe..6a62f60e558b 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/dense_linear_quantized_f16_wmma.loom @@ -1,19 +1,7 @@ -// Dense raw-quantized projection for the Qwen attention path. -// -// Each two-wave workgroup computes 64 output channels for 32 contiguous -// tokens. Raw GGUF Q4_K or Q6_K rows are decoded directly into a padded FP16 -// LDS tile, while the same workgroup converts the corresponding F32 activation -// tile to FP16. Four wave64 WMMA accumulators cover the 32x32 result owned by -// each wave. A wave-private LDS slice transposes each accumulator for coalesced -// F32 publication. -// -// Both weight contracts use the unmodified GGUF layout: -// Q4_K: [output channel][input size / 256][144 bytes] -// Q6_K: [output channel][input size / 256][210 bytes] -// The entry point selects the format before the shared device template is -// instantiated, so inactive decode logic is absent from the emitted kernel. No -// persistent repacking or expanded-weight allocation is required. -amdgpu.target @qwen3_moe_dense_gfx11_wave64 {subgroup_size = 64} +// Q4_K x Q8_1 projection kernels for the Qwen attention path. +template.decl @qwen3_moe.dense.q4k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) amdgpu.target @qwen3_moe_dense_gfx11_wave32 {subgroup_size = 32} @@ -29,190 +17,13 @@ config.decl @qwen3_moe.dense_quantized.output_accumulation : %value: index where config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] -// These declarations are used only by the differential correctness case. The -// dense production kernel has no routing input or dependency. -kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) - -kernel.decl @qwen3_moe_routed_linear_q4k_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) - // Shared raw-layout Q4_K row contraction primitive. Each lane consumes both // nibbles of one packed-code load. func.decl @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %lane: index) -> (f32) -// Shared raw Q6_K decoder linked from the GGML physical-layout module. -func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) - -// These declarations are used only by the Q6_K differential correctness case. -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.decl @ggml_linear_q6k_q8_1_x4(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %q8_input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count$63: index, %input_size$64: index) launch(%token_count$65: index, %input_size$66: index, %input: buffer, %output: buffer) -kernel.decl @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %rms_weight: buffer, %normalized_f32_output: buffer, %q8_1_x4_output: buffer) - -// Decodes the four adjacent Q4_K values owned by one load packet. One packed -// header load supplies both half scales and all twelve six-bit group scales; -// one packed code load supplies four adjacent unsigned nibbles. The final -// affine conversion is formed in F32 before truncation to the FP16 WMMA input. -func.def inline @qwen3_moe_dense_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %block_bytes = index.constant 144 : offset - %scale_offset = index.constant 4 : offset - %code_offset = index.constant 16 : offset - %c2_i32 = scalar.constant 2 : i32 - %c4_i32 = scalar.constant 4 : i32 - %c15 = vector.constant 15 : vector<1xi32> - %c48 = vector.constant 48 : vector<1xi32> - %q4_mask = vector.constant 252645135 : vector<1xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_offset : offset - %code_byte_base = index.add %block_byte_base, %code_offset : offset - %dm_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<3xi32> - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> - %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> - %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 - %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %q_page0 = index.div %bounded_group, %c2 : index - %q_page = index.mul %q_page0, %c8 : index - %q_word_index0 = index.add %q_page, %bounded_packet : index - %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index - %is_low = index.cmp ult, %bounded_group, %c4 : index - %scale_lane = index.rem %bounded_group, %c4 : index - %scale_shift_index = index.mul %scale_lane, %c8 : index - %scale_shift_i32 = index.cast %scale_shift_index : index to i32 - %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> - %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 - %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 - %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 - %scale0 = vector.splat %scale0_i32 : vector<1xi32> - %scale1 = vector.splat %scale1_i32 : vector<1xi32> - %scale2 = vector.splat %scale2_i32 : vector<1xi32> - %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 - %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 - %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> - %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> - %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 - %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 - %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> - %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> - %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> - %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> - %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> - %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> - %scale = vector.ori %scale_low, %scale_high : vector<1xi32> - %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> - %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> - %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> - %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> - %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> - %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> - %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> - %d_vector1 = vector.splat %d : vector<1xf32> - %dmin_vector1 = vector.splat %dmin : vector<1xf32> - %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> - %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> - %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 - %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 - %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> - %q_half = index.rem %bounded_group, %c2 : index - %q_shift_index = index.mul %q_half, %c4 : index - %q_shift_i32 = index.cast %q_shift_index : index to i32 - %q_shift = vector.splat %q_shift_i32 : vector<1xi32> - %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> - %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> - %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> - %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> - %negative_minimum_scale = scalar.negf %minimum_scale : f32 - %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 - %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 - %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 - %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 - %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 - %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 - %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 - %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 - %half0 = scalar.fptrunc %value0 : f32 to f16 - %half1 = scalar.fptrunc %value1 : f32 to f16 - %half2 = scalar.fptrunc %value2 : f32 to f16 - %half3 = scalar.fptrunc %value3 : f32 to f16 - %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> - func.return %result : vector<4xf16> -} - -// Low-fixed-cost Q4_K body shared by the ordinary and completion-fused decode -// exports. Each wave owns one output channel and contracts a Q8_1 activation -// row directly against the original GGUF weight row. -func.template device @qwen3_moe_dense_linear_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.dense_quantized.input_size : index - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144_bytes = index.constant 144 : offset - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %token, %launch_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %q4_block_count, %c144_bytes : index, offset -> offset - %weight_row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_row_byte_base = index.scale %token, %q8_row_bytes : index, offset -> offset - %q8_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %weight, %output : buffer, buffer, buffer - %lane_sum = scf.if %publish_output -> (f32) { - %channel_sum = scf.if %valid_channel -> (f32) { - %sum = func.call @qwen3_moe_q4k_q8_1_x4_paired_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - scf.yield %channel_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %dot = kernel.subgroup.reduce %lane_sum : f32 - %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index - scf.if %publish_output { - scf.if %valid_channel { - scf.if %is_lane_zero { - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%bounded_output_size]xf32> - %result = scf.if %accumulates_output -> (f32) { - %residual = view.load %output_view[%token, %channel] : view<[%launch_token_count]x[%bounded_output_size]xf32> -> f32 - %sum = scalar.addf %residual, %dot : f32 - scf.yield %sum : f32 - } else { - scf.yield %dot : f32 - } - view.store %result, %output_view[%token, %channel] : f32, view<[%launch_token_count]x[%bounded_output_size]xf32> - } - } - } - func.return -} +kernel.decl @qwen3_moe_rmsnorm_f32_quantize_q8_1_x4(%token_count$78: index) launch(%token_count$79: index, %input: buffer, %rms_weight: buffer, %normalized_f32_output: buffer, %q8_1_x4_output: buffer) // Direct decode entry point. This complements the grouped WMMA schedule below: // it gives up cross-token weight reuse in exchange for doing no padded matrix @@ -234,449 +45,11 @@ kernel.def target(@qwen3_moe_dense_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_ %c0 = index.constant 0 : index %valid_token = index.cmp ult, %token0, %bounded_token_count : index %safe_token = scf.select %valid_token, %token0, %c0 : index - func.apply(%valid_token, %bounded_token_count, %safe_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) + template.apply<@qwen3_moe.dense.q4k_q8_1_x4.body>(%valid_token, %bounded_token_count, %safe_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) kernel.return } -// Decode-only route that publishes the normalized F32 and Q8_1 x4 rows -// consumed by the following feed-forward boundary. -kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count: index) { - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c256 = index.constant 256 : index - %output_tiles = index.div %output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %publish_output = scalar.constant true : i1 - %projection_token = kernel.workgroup.id : index - func.apply(%publish_output, %bounded_token_count, %projection_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) - %output_size0 = config.get @qwen3_moe.dense_quantized.output_size : index - %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index - %token0 = kernel.workgroup.id : index - %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index - %c0 = index.constant 0 : index - %c8 = index.constant 8 : index - %workitem = kernel.workitem.id : index - %output_tile_count = index.div %output_size, %c8 : index - %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %normalized_output, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - // Publish every producer's residual stores before the leader advances one - // workgroup arrival. The last arrival then acquires the complete row. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %output_tile_count_i32 = index.cast %output_tile_count : index to i32 - %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 - %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 - %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 - scf.if %is_last_output_tile { - kernel.barrier scope(workgroup) ordering(acquire) - %publish_normalized = scalar.constant true : i1 - func.apply(%publish_normalized, %c8, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) - // Reset only after every normalized F32 and Q8 store completes. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> - } - } - kernel.return -} - -// Shared matrix-tile schedule. Callers own the launch grid and pass one -// validated tile coordinate plus a literal storage kind, allowing the linker -// and JIT to erase the inactive packed decoder. -func.template device @qwen3_moe_dense_linear_quantized_f16_wmma_body(%weight_format: index, %token_count: index, %input_size0: index, %output_size0: index, %output_accumulation: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %output: 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, 256)] : 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 - %c6 = index.constant 6 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c40 = index.constant 40 : index - %c64 = index.constant 64 : index - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %q4_block_bytes = index.constant 144 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 1024 : offset - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %zero_accumulator = vector.constant 0.0 : vector<8xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %is_q6 = index.cmp eq, %weight_format, %c6 : index - %weight_block_bytes = scf.select %is_q6, %c210_bytes, %q4_block_bytes : offset - %quant_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %quant_block_count, %weight_block_bytes : index, offset -> offset - %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_output_size]xf32> - %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 = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group = [%c0 to %c8 step %c1](%acc00 = %block_acc00 : vector<8xf16>, %acc01 = %block_acc01 : vector<8xf16>, %acc10 = %block_acc10 : vector<8xf16>, %acc11 = %block_acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %block_k_base = index.mul %quant_block, %c256 : index - %group_k_add = index.mul %quant_group, %c32 : index - %k_origin = index.add %block_k_base, %group_k_add : index - scf.for %row_offset = [%c0 to %c64 step %c16] unroll { - %local_row0 = index.add %load_row, %row_offset : index - %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %weight_values = scf.if %valid_channel -> (vector<4xf16>) { - %row_byte_base = index.scale %channel, %weight_row_bytes : index, offset -> offset - %decoded = scf.if %is_q6 -> (vector<4xf16>) { - %q6_values = func.call @ggml_q6k_f16_vector4(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) - scf.yield %q6_values : vector<4xf16> - } else { - %q4_values = func.call @qwen3_moe_dense_q4k_wmma_vector4(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) - scf.yield %q4_values : 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 - %activation_values = scf.if %valid_token -> (vector<4xf16>) { - %bounded_token, %input_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> - scf.yield %converted : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> - } - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { - %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - // WMMA owns [channel][token] fragments. Each wave transposes one fragment - // through its private LDS slice so every lane publishes four contiguous F32 - // channels for one logical token. - %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 - %accumulates_output = index.cmp eq, %output_accumulation, %c1 : index - 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 { - %bounded_token, %output_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - %published = scf.if %accumulates_output -> (vector<4xf32>) { - %residual = vector.load.mask %output_view[%bounded_token, %channel0], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> - %sum = vector.addf %residual, %wide : vector<4xf32> - scf.yield %sum : vector<4xf32> - } else { - scf.yield %wide : vector<4xf32> - } - vector.store.mask %published, %output_view[%bounded_token, %channel0], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - 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 { - %bounded_token, %output_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - %published = scf.if %accumulates_output -> (vector<4xf32>) { - %residual = vector.load.mask %output_view[%bounded_token, %channel0], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> - %sum = vector.addf %residual, %wide : vector<4xf32> - scf.yield %sum : vector<4xf32> - } else { - scf.yield %wide : vector<4xf32> - } - vector.store.mask %published, %output_view[%bounded_token, %channel0], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - 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 { - %bounded_token, %output_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - %published = scf.if %accumulates_output -> (vector<4xf32>) { - %residual = vector.load.mask %output_view[%bounded_token, %channel1], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> - %sum = vector.addf %residual, %wide : vector<4xf32> - scf.yield %sum : vector<4xf32> - } else { - scf.yield %wide : vector<4xf32> - } - vector.store.mask %published, %output_view[%bounded_token, %channel1], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - 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 { - %bounded_token, %output_token_count = index.assume %token1, %bounded_token_count [lt(%token1, %bounded_token_count)] : index, index - %values = vector.load %result_physical_view[%publish_token, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - %published = scf.if %accumulates_output -> (vector<4xf32>) { - %residual = vector.load.mask %output_view[%bounded_token, %channel1], %mask, %c0_f32x4 : view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> - %sum = vector.addf %residual, %wide : vector<4xf32> - scf.yield %sum : vector<4xf32> - } else { - scf.yield %wide : vector<4xf32> - } - vector.store.mask %published, %output_view[%bounded_token, %channel1], %mask : vector<4xf32>, view<[%bounded_token_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - func.return -} - -// Shared launch geometry for configured standalone kernels and parameterized -// command-program kernels. -func.template @qwen3_moe_dense_linear_quantized_f16_wmma_launch(%token_capacity: index, %output_size: index) -> (index, index, index, 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 - func.return %output_tiles, %token_tiles, %c1, %c128 : index, index, index, index -} - -// Configured entry points retain the compact standalone kernel ABI used by -// library-style callers. Shape configs specialize the native code once while -// token count remains the only per-dispatch scalar. -kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q4k_f16_wmma(%token_count: index) { - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %output_tiles, %token_tiles, %c1, %c128 = func.apply(%token_capacity, %output_size) : (index, index) -> (index, index, index, 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 @qwen3_moe.dense_quantized.input_size : index - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %q4 = index.constant 4 : index - func.apply(%q4, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q6k_f16_wmma(%token_count: index) { - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %output_tiles, %token_tiles, %c1, %c128 = func.apply(%token_capacity, %output_size) : (index, index) -> (index, index, index, 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 @qwen3_moe.dense_quantized.input_size : index - %output_size = config.get @qwen3_moe.dense_quantized.output_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %output_accumulation = config.get @qwen3_moe.dense_quantized.output_accumulation : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %q6 = index.constant 6 : index - func.apply(%q6, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -// Parameterized entry points carry launch-varying shapes through command IR. -// Command planning materializes each exact fact environment in a private unit, -// so these scalar values do not survive in the native device ABI. -kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q4k_f16_wmma_parameterized(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index) { - %output_tiles, %token_tiles, %c1, %c128 = func.apply(%token_count, %output_size) : (index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048), range(%input_size, 256, 32768), mul(%input_size, 256), range(%output_size, 1, 262144), range(%output_accumulation, 0, 1)] { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %q4 = index.constant 4 : index - func.apply(%q4, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_dense_gfx11_wave64) @qwen3_moe_dense_linear_q6k_f16_wmma_parameterized(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index) { - %output_tiles, %token_tiles, %c1, %c128 = func.apply(%token_count, %output_size) : (index, index) -> (index, index, index, index) - kernel.launch.config workgroups(%output_tiles, %token_tiles, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %output_size: index, %output_accumulation: index, %input: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048), range(%input_size, 256, 32768), mul(%input_size, 256), range(%output_size, 1, 262144), range(%output_accumulation, 0, 1)] { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %channel_tile = kernel.workgroup.id : index - %token_tile = kernel.workgroup.id : index - %q6 = index.constant 6 : index - func.apply(%q6, %bounded_token_count, %input_size, %output_size, %output_accumulation, %channel_tile, %token_tile, %input, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer) - kernel.return -} - -// The exact-representable activation and nonzero packed weight bytes compare -// dense addressing against the already-certified routed WMMA provider. The -// shape crosses both the 32-token and 64-channel tile boundaries. -check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_differential_case { - %token_count = check.literal value(33) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(1) : index - %expert_count = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<33x512xf32> - %route_ids = check.generate.fill value(0) : tensor<33x1xi32> - %expert_table = check.generate.fill value(-1) : tensor<34xi32> - %weight = check.generate.fill value(34) : tensor<1x65x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<33x65xf32> - %actual = check.generate.fill value(1.0) : tensor<33x65xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<33x1xi32>, tensor<34xi32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %expected) : [index](index, tensor<33x512xf32>, tensor<34xi32>, tensor<1x65x2x144xi8>, tensor<33x65xf32>) - kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<33x512xf32>, tensor<1x65x2x144xi8>, tensor<33x65xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<33x65xf32> - check.return -} - -// Q6_K uses the same exact-representable activation across both the Q8_1 dot -// reference and FP16 WMMA provider. The comparison crosses both matrix tile -// boundaries and validates the format-specialized packed decoder. -check.case public @qwen3_moe_dense_linear_q6k_f16_wmma_differential_case { - %token_count = check.literal value(33) : index - %input_size = check.literal value(512) : index - %output_size = check.literal value(65) : index - %input = check.generate.fill value(0.00390625) : tensor<33x512xf32> - %q8_input = check.generate.fill value(0) : tensor<33x576xi8> - %weight = check.generate.fill value(-86) : tensor<65x2x210xi8> - %expected = check.generate.fill value(0.0) : tensor<33x65xf32> - %actual = check.generate.fill value(1.0) : tensor<33x65xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<33x512xf32>, tensor<33x576xi8>) - kernel.launch @ggml_linear_q6k_q8_1_x4[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %q8_input, %weight, %expected) : [index, index, index](index, index, index, tensor<33x576xi8>, tensor<65x2x210xi8>, tensor<33x65xf32>) - kernel.launch @qwen3_moe_dense_linear_q6k_f16_wmma[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<33x512xf32>, tensor<65x2x210xi8>, tensor<33x65xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<33x65xf32> - check.return -} - -// The O-projection dimensions double the Q/K input depth. This nonzero -// production-size differential keeps that loop regime covered independently -// of the benchmark parameter sweep. -check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_o_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(4096) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(1) : index - %expert_count = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<1x4096xf32> - %q8_input = check.generate.fill value(0) : tensor<1x4608xi8> - %route_ids = check.generate.fill value(0) : tensor<1x1xi32> - %expert_table = check.generate.fill value(-1) : tensor<2xi32> - %weight = check.generate.fill value(34) : tensor<1x2048x16x144xi8> - %expected = check.generate.fill value(0.0) : tensor<1x2048xf32> - %actual = check.generate.fill value(0.0) : tensor<1x2048xf32> - %direct = check.generate.fill value(0.0) : tensor<1x2048xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<1x1xi32>, tensor<2xi32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x4096xf32>, tensor<1x4608xi8>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %weight, %expected) : [index](index, tensor<1x4096xf32>, tensor<2xi32>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %actual) : [index](index, tensor<1x4096xf32>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %direct) : [index](index, tensor<1x4608xi8>, tensor<1x2048x16x144xi8>, tensor<1x2048xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<1x2048xf32> - check.expect.close actual(%direct) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<1x2048xf32> - check.return -} - -// Crosses the direct schedule's eight-wave output tile while comparing its -// Q8_1 dot path against the independently structured FP16 WMMA provider. -check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_differential_case { - %token_count = check.literal value(2) : index - %input_size = check.literal value(512) : index - %input = check.generate.fill value(0.00390625) : tensor<2x512xf32> - %q8_input = check.generate.fill value(0) : tensor<2x576xi8> - %weight = check.generate.iota offset(-72) step(1) period(144) : tensor<65x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<2x65xf32> - %actual = check.generate.fill value(0.0) : tensor<2x65xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<2x512xf32>, tensor<2x576xi8>) - kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %expected) : [index](index, tensor<2x512xf32>, tensor<65x2x144xi8>, tensor<2x65xf32>) - kernel.launch @qwen3_moe_dense_linear_q4k_q8_1_x4[%token_count](%token_count, %q8_input, %weight, %actual) : [index](index, tensor<2x576xi8>, tensor<65x2x144xi8>, tensor<2x65xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<2x65xf32> - check.return -} +kernel.decl @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count$1: index) launch(%token_count$2: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) // A zero projection leaves the nonzero caller-provided residual unchanged. // Overwrite publication would instead produce zero, making the accumulation @@ -758,17 +131,6 @@ check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_benchmark check.return } -check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case { - %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %weight = check.generate.fill value(0) : tensor<4096x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x4096xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> - kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<4096x8x144xi8>, tensor<[%token_count]x4096xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x4096xf32> - check.return -} - check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case { %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index %input_size = check.literal value(2048) : index @@ -783,17 +145,6 @@ check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_cas check.return } -check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case { - %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %weight = check.generate.fill value(0) : tensor<512x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> - kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<512x8x144xi8>, tensor<[%token_count]x512xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> - check.return -} - check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case { %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index %input_size = check.literal value(2048) : index @@ -808,28 +159,6 @@ check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_cas check.return } -check.case public @qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case { - %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %weight = check.generate.fill value(0) : tensor<512x8x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x512xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x512xf32> - kernel.launch @qwen3_moe_dense_linear_q6k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<512x8x210xi8>, tensor<[%token_count]x512xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x512xf32> - check.return -} - -check.case public @qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case { - %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x4096xf32> - %weight = check.generate.fill value(0) : tensor<2048x16x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_dense_linear_q4k_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<[%token_count]x4096xf32>, tensor<2048x16x144xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case { %token_count = check.param.choice values([1, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index %input_size = check.literal value(4096) : index @@ -844,26 +173,10 @@ check.case public @qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_cas check.return } -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_differential_case> @qwen3_moe_dense_linear_q4k_f16_wmma_differential - -check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_differential_case> @qwen3_moe_dense_linear_q6k_f16_wmma_differential - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_differential_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_differential - -check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_differential_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_differential - check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_decode check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8_composed_decode -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_decode {token_count = 1} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_q_prefill_512 {token_count = 512} - check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_decode {token_count = 1} check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_prefill_32 {token_count = 32} @@ -872,14 +185,6 @@ check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_q_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_q_prefill_512 {token_count = 512} -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_decode {token_count = 1} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_k_prefill_512 {token_count = 512} - check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_decode {token_count = 1} check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_prefill_32 {token_count = 32} @@ -888,22 +193,6 @@ check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_k_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_k_prefill_512 {token_count = 512} -check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_decode {token_count = 1} - -check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_dense_linear_q6k_f16_wmma_v_projection_benchmark_case> @qwen3_moe_dense_linear_q6k_f16_wmma_v_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_decode {token_count = 1} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_dense_linear_q4k_f16_wmma_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_f16_wmma_o_prefill_512 {token_count = 512} - check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_decode {token_count = 1} check.benchmark<@qwen3_moe_dense_linear_q4k_q8_1_x4_o_projection_benchmark_case> @qwen3_moe_dense_linear_q4k_q8_1_x4_o_prefill_32 {token_count = 32} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom index 8c912bbfc954..2eaba9e6d6f3 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/expert_table_partition_fused.loom @@ -6,169 +6,13 @@ // The last workgroup acquires all counts, compacts them into deterministic // 32-row descriptors, and resets the counter so reusable command buffers can // issue the kernel again against the same storage. -amdgpu.target @qwen3_moe_expert_table_partition_gfx11_wave32 {subgroup_size = 32} +target.decl @qwen3_moe_expert_table_partition_gfx11_wave32 -kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) +kernel.decl @ggml_moe_build_expert_table(%token_count$0: index, %route_count$1: index, %route_stride$2: index, %expert_count$3: index) launch(%token_count$4: index, %route_count$5: index, %route_stride$6: index, %expert_count$7: index, %route_ids: buffer, %expert_table: buffer) -kernel.decl @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) +kernel.decl @ggml_moe_build_expert_partition_table(%token_count$10: index, %route_count$11: index, %expert_count$12: index) launch(%token_count$13: index, %route_count$14: index, %expert_count$15: index, %expert_table: buffer, %partition_table: buffer) -kernel.def target(@qwen3_moe_expert_table_partition_gfx11_wave32) @qwen3_moe_build_expert_table_partition_prefill_512(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { - %c1 = index.constant 1 : index - %launch_expert_count = index.constant 128 : index - %workgroup_size = index.constant 256 : index - kernel.launch.config workgroups(%launch_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 512, 512)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 8, 8)] : index - %bounded_route_stride = index.assume %route_stride [range(%route_stride, 8, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 128, 128)] : index - %expert0 = kernel.workgroup.id : index - %expert, %launch_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index - %lane0 = kernel.workitem.id : index - %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %workgroup_size = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c7_i32 = scalar.constant 7 : i32 - %c13_i32 = scalar.constant 13 : i32 - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index - %assignment_table_byte_base = index.scale %launch_expert_count, %c4_bytes : index, offset -> offset - %route_ids_noalias, %expert_table_noalias, %partition_table_noalias, %completion_counter_noalias = buffer.assume.noalias %route_ids, %expert_table, %partition_table, %completion_counter : buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%launch_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%launch_expert_count]x[%bounded_token_count]xi32> - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> - - %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { - %assignment = index.add %block_base, %lane : index - %in_range = index.cmp ult, %assignment, %assignment_count : index - %route_expert_i32 = scf.if %in_range -> (i32) { - %token0 = index.div %assignment, %bounded_route_count : index - %route0 = index.rem %assignment, %bounded_route_count : index - %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index - %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - scf.yield %loaded : i32 - } else { - %cn1_i32 = scalar.constant -1 : i32 - scf.yield %cn1_i32 : i32 - } - %route_expert0 = index.cast %route_expert_i32 : i32 to index - %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 127)] : index - %matches = index.cmp eq, %route_expert, %expert : index - %match_i32 = scf.if %matches -> (i32) { - scf.yield %c1_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 - %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 - %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 - scf.if %matches { - %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 - %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index - %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 511)] : index - %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index - %assignment_i32 = index.cast %assignment : index to i32 - view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%launch_expert_count]x[%bounded_token_count]xi32> - } - %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 - scf.yield %next_matched_base : i32 - } - - %is_lane_zero = index.cmp eq, %lane, %c0 : index - scf.if %is_lane_zero { - view.store %expert_route_count, %count_view[%expert] : i32, view<[%launch_expert_count]xi32> - } - - // Every count store precedes its workgroup's release. The last arrival - // acquires all preceding releases before any lane loads the complete table. - %local_old_counter = scf.if %is_lane_zero -> (i32) { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 - scf.yield %old_counter : i32 - } else { - scf.yield %c0_i32 : i32 - } - %old_counter = kernel.workgroup.reduce %local_old_counter : i32 - %expert_count_i32 = index.cast %launch_expert_count : index to i32 - %last_expert_i32 = scalar.subi %expert_count_i32, %c1_i32 : i32 - %negative_expert_count_i32 = scalar.subi %c0_i32, %expert_count_i32 : i32 - %is_last_expert = scalar.cmpi eq, %old_counter, %last_expert_i32 : i32 - scf.if %is_last_expert { - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %maximum_partition_count = index.add %assignment_partition_count, %launch_expert_count : index - %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> - %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> - %has_expert = index.cmp ult, %lane, %launch_expert_count : index - %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { - %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index - %loaded = view.load %count_view[%partition_expert] : view<[%launch_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index - %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 512)] : index - %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index - %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index - %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 - %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 - %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 - %partition_count = index.cast %partition_count_i32 : i32 to index - %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index - %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index - %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 255)] : index - scf.if %has_expert { - %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index - %expert_i32 = index.cast %partition_expert : index to i32 - scf.for %partition = [%c0 to %expert_partition_count step %c1] { - %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index - %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index - %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index - %partition_remainder = index.rem %expert_assignment_count, %c32 : index - %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index - %partition_row_count = scf.if %has_partial_tail -> (index) { - %next_partition = index.add %partition, %c1 : index - %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index - %tail_row_count = scf.if %is_tail_partition -> (index) { - scf.yield %partition_remainder : index - } else { - scf.yield %c32 : index - } - scf.yield %tail_row_count : index - } else { - scf.yield %c32 : index - } - %partition_i32 = index.cast %partition : index to i32 - %partition_row_count_i32 = index.cast %partition_row_count : index to i32 - %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 - %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 - %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 - %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 - %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 - view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> - } - } - scf.if %is_lane_zero { - view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> - } - // The counter cannot become reusable until every descriptor store is - // complete and visible to the following dispatch. - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.if %is_lane_zero { - view.atomic.reduce %negative_expert_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> - } - } - kernel.return -} +kernel.decl @qwen3_moe_build_expert_table_partition_prefill_512(%token_count$20: index, %route_count$21: index, %route_stride$22: index, %expert_count$23: index) launch(%token_count$24: index, %route_count$25: index, %route_stride$26: index, %expert_count$27: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) // Compare both fused outputs against the production composition, then invoke // the fused route twice against one counter to make reset correctness visible. @@ -186,8 +30,8 @@ check.case public @qwen3_moe_expert_table_partition_fused_differential_case { %actual_partition_table1 = check.generate.fill value(-1) : tensor<257xi32> %completion_counter = check.generate.fill value(0) : tensor<1xi32> %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_expert_table) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expected_expert_table, %expected_partition_table) : [index, index, index](index, index, index, tensor<65664xi32>, tensor<257xi32>) + kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expected_expert_table) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>) + kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expected_expert_table, %expected_partition_table) : [index, index, index](index, index, index, tensor<65664xi32>, tensor<257xi32>) kernel.launch @qwen3_moe_build_expert_table_partition_prefill_512[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_expert_table0, %actual_partition_table0, %completion_counter) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>, tensor<257xi32>, tensor<1xi32>) kernel.launch @qwen3_moe_build_expert_table_partition_prefill_512[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %actual_expert_table1, %actual_partition_table1, %completion_counter) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>, tensor<257xi32>, tensor<1xi32>) check.expect.equal actual(%actual_expert_table0) expected(%expected_expert_table) : tensor<65664xi32> @@ -219,11 +63,13 @@ check.case public @qwen3_moe_expert_table_partition_composed_benchmark_case { %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<512x8xi32> %expert_table = check.generate.fill value(-1) : tensor<65664xi32> %partition_table = check.generate.fill value(-1) : tensor<257xi32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<65664xi32>, tensor<257xi32>) + kernel.launch @ggml_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<512x8xi32>, tensor<65664xi32>) + kernel.launch @ggml_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<65664xi32>, tensor<257xi32>) check.return } +check.benchmark<@qwen3_moe_expert_table_partition_fused_differential_case> @qwen3_moe_expert_table_partition_fused_differential + check.benchmark<@qwen3_moe_expert_table_partition_fused_benchmark_case> @qwen3_moe_expert_table_partition_fused_prefill_512 check.benchmark<@qwen3_moe_expert_table_partition_composed_benchmark_case> @qwen3_moe_expert_table_partition_composed_prefill_512 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom deleted file mode 100644 index fc523520717a..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_f32_f16_wmma.loom +++ /dev/null @@ -1,690 +0,0 @@ -// Qwen3 MoE grouped-query decode FlashAttention. -// -// One or two four-wave workgroups compute up to 16 query heads that share one -// KV head. Each workgroup walks the exact KV length in 64-row blocks while -// carrying online-softmax max, sum, and output state in registers. With two -// output partitions, each workgroup owns one disjoint 64-channel output half; -// this duplicates QK work but exposes more parallelism without publishing -// partial softmax state. The ownership changes mirror the cooperative-matrix -// schedule used by llama.cpp's Vulkan CM1 kernel: -// -// 1. All workitems stage a scaled 16x128 F16 GQA-head tile. -// 2. Each wave computes one 16x16 QK score slice. -// 3. Scores cross LDS so each wave can normalize four complete query heads. -// 4. F16 probabilities cross LDS for four P*V WMMA steps. -// 5. Each active lane retains one four-channel F16 packet for each of its -// four query heads across subsequent KV blocks. -// -// K and V remain in the row-major llama.cpp cache layout -// [KV token][KV head][128]. Their aligned F16 fragments load directly from -// global memory; there is no expanded or repacked persistent allocation. QK -// and the online-softmax statistics remain F32, while the P*V accumulation and -// carried output match the Vulkan oracle's F16 policy. Unlike the split-K -// fallback, this kernel never publishes partial tensors or completion atomics. -amdgpu.target @qwen3_moe_attention_decode_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] - -config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] - -// Number of independent 64-channel output partitions per KV head. One avoids -// redundant QK work; two exposes more workgroups for short decode contexts. -config.decl @qwen3_moe.attention.decode.output_partition_count : %value: index where [range(%value, 1, 2)] - -kernel.def target(@qwen3_moe_attention_decode_gfx11_wave64) @qwen3_moe_flash_attention_decode_f32_f16_wmma(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %output_partition_count = config.get @qwen3_moe.attention.decode.output_partition_count : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %workgroup_count = index.mul %key_value_head_count, %output_partition_count : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) { - %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %output_partition_count = config.get @qwen3_moe.attention.decode.output_partition_count : index - %workgroup_x0 = kernel.workgroup.id : index - %workgroup_count0 = index.mul %key_value_head_count, %output_partition_count : index - %workgroup_x, %workgroup_count = index.assume %workgroup_x0, %workgroup_count0 [lt(%workgroup_x0, %workgroup_count0)] : index, index - %key_value_head0 = index.div %workgroup_x, %output_partition_count : index - %key_value_head = index.assume %key_value_head0 [range(%key_value_head0, 0, 63)] : index - %output_partition = index.rem %workgroup_x, %output_partition_count : 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 - %c2 = index.constant 2 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c15 = index.constant 15 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %query_stage_bytes = index.constant 4352 : offset - %score_stage_bytes = index.constant 6144 : offset - %probability_stage_bytes = index.constant 3072 : offset - %product_stage_bytes = index.constant 2048 : offset - %tail_key_value_stage_capacity = index.constant 8192 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %head_size_f32 = scalar.constant 128.0 : f32 - %attention_scale = scalar.rsqrtf %head_size_f32 : f32 - %c0_f16 = scalar.constant 0.0 : f16 - %c1_f32 = scalar.constant 1.0 : f32 - %output_zero0 = vector.constant 0.0 : vector<4xf16> - %output_zero1 = vector.constant 0.0 : vector<4xf16> - %output_zero2 = vector.constant 0.0 : vector<4xf16> - %output_zero3 = vector.constant 0.0 : vector<4xf16> - %c0_f16x8 = vector.constant 0.0 : vector<8xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %key_value_width = index.mul %key_value_head_count, %c128 : index - %key_value_head_base = index.mul %key_value_head, %c128 : index - %full_key_value_block_count = index.div %bounded_key_value_token_count, %c64 : index - %full_key_value_token_count0 = index.mul %full_key_value_block_count, %c64 : index - %full_key_value_token_count = index.assume %full_key_value_token_count0 [range(%full_key_value_token_count0, 0, 32768), mul(%full_key_value_token_count0, 64)] : index - %tail_key_value_token_count = index.sub %bounded_key_value_token_count, %full_key_value_token_count : index - %has_key_value_tail = index.cmp ne, %tail_key_value_token_count, %c0 : index - %has_single_key_value_tail = index.cmp eq, %tail_key_value_token_count, %c1 : index - %tail_key_value_stage_bytes = scf.select %has_key_value_tail, %tail_key_value_stage_capacity, %c0_offset : offset - %subgroup_score_column = index.mul %subgroup, %c16 : index - %subgroup_query_row = index.mul %subgroup, %c4 : index - %query_row0 = index.add %subgroup_query_row, %c0 : index - %query_row1 = index.add %subgroup_query_row, %c1 : index - %query_row2 = index.add %subgroup_query_row, %c2 : index - %query_row3 = index.add %subgroup_query_row, %c3 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head_base, %query_row1 : index - %query_head2 = index.add %query_head_base, %query_row2 : index - %query_head3 = index.add %query_head_base, %query_row3 : index - %query_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_valid1 = index.cmp ult, %query_head1, %query_head_count : index - %query_valid2 = index.cmp ult, %query_head2, %query_head_count : index - %query_valid3 = index.cmp ult, %query_head3, %query_head_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 - %output_tile_count = index.div %c2, %output_partition_count : index - %output_tile_end = index.add %output_partition, %output_tile_count : index - %lane_product_channel0 = index.rem %lane, %c16 : index - %lane_product_channel = index.mul %lane_product_channel0, %c4 : index - %lane_output_channel = index.mul %lane, %c4 : index - %lane_output_at_or_after_partition = index.cmp uge, %lane_output_tile, %output_partition : index - %lane_output_before_end = index.cmp ult, %lane_output_tile, %output_tile_end : index - %lane_has_output = scalar.andi %lane_output_at_or_after_partition, %lane_output_before_end : i1 - // The padded LDS rows mirror the Vulkan oracle's ownership changes. Q uses - // eight spare F16 columns after its 128 channels. Score and probability - // transpose to key-major rows with eight spare columns after 16 queries. - // These strides avoid the bank pattern produced by dense transposed rows. - %query_transposed_layout = encoding.layout.strided [1, 136] : encoding - %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding - %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer - %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer - %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer - %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer - %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %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 - %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer - %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer - %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> - %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> - %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> - %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> - %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x128xf16> - // Scale and truncate Q exactly once. The Vulkan reference does this before - // entering its KV loop, making QK a native F16 WMMA while retaining F32 - // accumulation. - scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %local_query_row = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %local_query_head = index.add %query_head_base, %local_query_row : index - %query_valid_load = index.cmp ult, %local_query_head, %query_head_count : index - %query_value = scf.if %query_valid_load -> (f16) { - %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 - %scaled = scalar.mulf %loaded, %attention_scale : f32 - %truncated = scalar.fptrunc %scaled : f32 to f16 - scf.yield %truncated : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> - } - 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>) { - // Four independent wave-level WMMAs produce a 16x64 score tile. - %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index - %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index - %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index - %score_init_values = vector.constant 0.0 : vector<4xf32> - %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> - %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_channel = index.add %key_value_head_base, %head_tile : index - %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // LDS transposes ownership from one 16-column score slice per wave to - // four complete query rows per wave. Every lane contributes one key - // column to each of those rows. - %key_token0 = index.add %key_origin, %lane : index - %key_token = index.assume %key_token0 [lt(%key_token0, %bounded_key_value_token_count)] : index - %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 - %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 - %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 - %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 - %mask_f16 = view.load %mask_view[%key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - %masked_score0 = scf.if %query_valid0 -> (f32) { - %score = scalar.addf %raw_score0, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score1 = scf.if %query_valid1 -> (f32) { - %score = scalar.addf %raw_score1, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score2 = scf.if %query_valid2 -> (f32) { - %score = scalar.addf %raw_score2, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score3 = scf.if %query_valid3 -> (f32) { - %score = scalar.addf %raw_score3, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %old_delta = vector.subf %current_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Match the Vulkan CM1 ownership schedule by computing two sequential - // 64-channel output tiles. This keeps one P*V accumulator live per wave - // and reuses a 2 KiB exchange tile instead of retaining both halves. - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { - %tile_at_or_after_partition = index.cmp uge, %output_tile, %output_partition : index - %tile_before_partition_end = index.cmp ult, %output_tile, %output_tile_end : index - %tile_selected = scalar.andi %tile_at_or_after_partition, %tile_before_partition_end : i1 - %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %tile_selected -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel0 = index.add %key_value_head_base, %output_tile_channel : index - %value_channel = index.add %value_channel0, %subgroup_product_channel : index - %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll { - %value_token0 = index.add %key_origin, %key_tile : index - %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> - } - vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index - %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %next_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> - %next_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> - %next_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> - %next_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // Complete every read before another output tile overwrites LDS. - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : 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> - } - // A single trailing KV row is cheaper as a native wave reduction and direct - // V update. Skip the general WMMA tail loop for that exact JIT-specialized - // shape; larger tails continue through the masked 32-row schedule below. - %wmma_tail_start = scf.select %has_single_key_value_tail, %bounded_key_value_token_count, %full_key_value_token_count : index - %tail_score_wave = index.cmp ult, %subgroup, %c2 : index - %wmma_tail_max, %wmma_tail_sum, %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 = scf.for %tail_key_origin = [%wmma_tail_start to %bounded_key_value_token_count step %c32](%current_max = %full_max : vector<4xf32>, %current_sum = %full_sum : vector<4xf32>, %current_output0 = %full_output0 : vector<4xf16>, %current_output1 = %full_output1 : vector<4xf16>, %current_output2 = %full_output2 : vector<4xf16>, %current_output3 = %full_output3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %tail_remaining = index.sub %bounded_key_value_token_count, %tail_key_origin : index - %tail_key_count = index.min %tail_remaining, %c32 : index - // Cooperatively stage one K tile, explicitly zeroing the padded rows. - scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %tail_key_row = index.div %linear, %c128 : index - %tail_key_channel = index.rem %linear, %c128 : index - %tail_key_valid = index.cmp ult, %tail_key_row, %tail_key_count : index - %tail_key_value = scf.if %tail_key_valid -> (f16) { - %tail_key_token0 = index.add %tail_key_origin, %tail_key_row : index - %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index - %global_key_channel = index.add %key_value_head_base, %tail_key_channel : index - %loaded = view.load %key_view[%tail_key_token, %global_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %tail_key_value, %tail_key_value_stage_view[%tail_key_row, %tail_key_channel] : f16, view<32x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Waves zero and one compute the two 16-column QK fragments in this tile. - scf.if %tail_score_wave { - %tail_score_subgroup = index.assume %subgroup [range(%subgroup, 0, 1)] : index - %tail_score_column = index.mul %tail_score_subgroup, %c16 : index - %tail_score_init_values = vector.constant 0.0 : vector<4xf32> - %tail_score_init = vector.fragment %tail_score_init_values shape [%m, %n] : vector<4xf32> - %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %tail_score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_fragment = vector.fragment.load %tail_key_value_stage_view[%tail_score_column, %head_tile] shape [%m, %k] : view<32x128xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %tail_score_fragment, %score_stage_view[%tail_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // LDS changes ownership from the QK wave to four query-row waves. Lanes - // beyond the logical tail never read the score or mask buffers. - %tail_lane_valid = index.cmp ult, %lane, %tail_key_count : index - %tail_valid0 = scalar.andi %tail_lane_valid, %query_valid0 : i1 - %tail_valid1 = scalar.andi %tail_lane_valid, %query_valid1 : i1 - %tail_valid2 = scalar.andi %tail_lane_valid, %query_valid2 : i1 - %tail_valid3 = scalar.andi %tail_lane_valid, %query_valid3 : i1 - %tail_valid = vector.from_elements %tail_valid0, %tail_valid1, %tail_valid2, %tail_valid3 : vector<4xi1> - %tail_key_token0 = index.add %tail_key_origin, %lane : index - %tail_mask_f32 = scf.if %tail_lane_valid -> (f32) { - %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index - %mask_f16 = view.load %mask_view[%tail_key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - scf.yield %mask_f32 : f32 - } else { - scf.yield %c0_f32 : f32 - } - %masked_score0 = scf.if %tail_valid0 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score1 = scf.if %tail_valid1 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score2 = scf.if %tail_valid2 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score3 = scf.if %tail_valid3 -> (f32) { - %raw_score = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 - %score = scalar.addf %raw_score, %tail_mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %tail_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %old_delta = vector.subf %current_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Reuse product scratch for V after every probability is resident in its - // disjoint LDS tile. - scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %tail_value_row = index.div %linear, %c128 : index - %tail_value_channel = index.rem %linear, %c128 : index - %tail_value_valid = index.cmp ult, %tail_value_row, %tail_key_count : index - %tail_value = scf.if %tail_value_valid -> (f16) { - %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index - %tail_value_token = index.assume %tail_value_token0 [lt(%tail_value_token0, %bounded_key_value_token_count)] : index - %global_value_channel = index.add %key_value_head_base, %tail_value_channel : index - %loaded = view.load %value_view[%tail_value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %tail_value, %tail_key_value_stage_view[%tail_value_row, %tail_value_channel] : f16, view<32x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Keep tail V staging separate from the 2 KiB product exchange, then use - // the same two 64-channel phases as the aligned path. - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { - %tile_at_or_after_partition = index.cmp uge, %output_tile, %output_partition : index - %tile_before_partition_end = index.cmp ult, %output_tile, %output_tile_end : index - %tile_selected = scalar.andi %tile_at_or_after_partition, %tile_before_partition_end : i1 - %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %tile_selected -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index - %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x128xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> - } - vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index - %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %next_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> - %next_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> - %next_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> - %next_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next_tile_output0, %next_tile_output1, %next_tile_output2, %next_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : 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> - } - %final_max, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3 = scf.if %has_single_key_value_tail -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %single_tail_token = index.assume %full_key_value_token_count [lt(%full_key_value_token_count, %bounded_key_value_token_count)] : index - %single_tail_channel = index.mul %lane, %c2 : index - %single_tail_key_channel = index.add %key_value_head_base, %single_tail_channel : index - %key_pair_f16 = vector.load %key_view[%single_tail_token, %single_tail_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<2xf16> - %query_pair0_f16 = vector.load %query_stage_view[%query_row0, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %query_pair1_f16 = vector.load %query_stage_view[%query_row1, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %query_pair2_f16 = vector.load %query_stage_view[%query_row2, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %query_pair3_f16 = vector.load %query_stage_view[%query_row3, %single_tail_channel] : view<16x136xf16> -> vector<2xf16> - %key_pair = vector.extf %key_pair_f16 : vector<2xf16> to vector<2xf32> - %query_pair0 = vector.extf %query_pair0_f16 : vector<2xf16> to vector<2xf32> - %query_pair1 = vector.extf %query_pair1_f16 : vector<2xf16> to vector<2xf32> - %query_pair2 = vector.extf %query_pair2_f16 : vector<2xf16> to vector<2xf32> - %query_pair3 = vector.extf %query_pair3_f16 : vector<2xf16> to vector<2xf32> - %product_pair0 = vector.mulf %query_pair0, %key_pair : vector<2xf32> - %product_pair1 = vector.mulf %query_pair1, %key_pair : vector<2xf32> - %product_pair2 = vector.mulf %query_pair2, %key_pair : vector<2xf32> - %product_pair3 = vector.mulf %query_pair3, %key_pair : vector<2xf32> - %partial_score0 = vector.reduce %product_pair0, %c0_f32 : vector<2xf32>, f32 - %partial_score1 = vector.reduce %product_pair1, %c0_f32 : vector<2xf32>, f32 - %partial_score2 = vector.reduce %product_pair2, %c0_f32 : vector<2xf32>, f32 - %partial_score3 = vector.reduce %product_pair3, %c0_f32 : vector<2xf32>, f32 - %partial_scores = vector.from_elements %partial_score0, %partial_score1, %partial_score2, %partial_score3 : vector<4xf32> - %reduced_scores = kernel.subgroup.reduce %partial_scores : vector<4xf32> - %mask_f16 = view.load %mask_view[%single_tail_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - %mask_vector = vector.splat %mask_f32 : vector<4xf32> - %raw_scores = vector.addf %reduced_scores, %mask_vector : vector<4xf32> - %masked_scores = vector.select %query_valid, %raw_scores, %negative_f32x4 : vector<4xf32> - %next_max = vector.maxnumf %wmma_tail_max, %masked_scores : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %old_delta = vector.subf %wmma_tail_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %wmma_tail_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %probability : vector<4xf32> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.if %lane_has_output -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %output_packet_channel = index.assume %lane_output_channel [range(%lane_output_channel, 0, 124), mul(%lane_output_channel, 4)] : index - %value_channel = index.add %key_value_head_base, %output_packet_channel : index - %value_packet = vector.load %value_view[%single_tail_token, %value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<4xf16> - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %probability0_scalar = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1_scalar = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2_scalar = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3_scalar = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %probability0 = vector.splat %probability0_scalar : vector<4xf16> - %probability1 = vector.splat %probability1_scalar : vector<4xf16> - %probability2 = vector.splat %probability2_scalar : vector<4xf16> - %probability3 = vector.splat %probability3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %wmma_tail_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %wmma_tail_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %wmma_tail_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %wmma_tail_output3, %old_scale3 : vector<4xf16> - %tail_output0 = vector.mulf %value_packet, %probability0 : vector<4xf16> - %tail_output1 = vector.mulf %value_packet, %probability1 : vector<4xf16> - %tail_output2 = vector.mulf %value_packet, %probability2 : vector<4xf16> - %tail_output3 = vector.mulf %value_packet, %probability3 : vector<4xf16> - %updated_output0 = vector.addf %scaled_current_output0, %tail_output0 : vector<4xf16> - %updated_output1 = vector.addf %scaled_current_output1, %tail_output1 : vector<4xf16> - %updated_output2 = vector.addf %scaled_current_output2, %tail_output2 : vector<4xf16> - %updated_output3 = vector.addf %scaled_current_output3, %tail_output3 : vector<4xf16> - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %wmma_tail_max, %wmma_tail_sum, %wmma_tail_output0, %wmma_tail_output1, %wmma_tail_output2, %wmma_tail_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // Normalize and publish the lane-owned packets. Query-tail rows never - // participate in the output store. - scf.if %lane_has_output { - %output_packet_channel = index.assume %lane_output_channel [range(%lane_output_channel, 0, 124), mul(%lane_output_channel, 4)] : 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 - %sum3_scalar = vector.extract %final_sum[3] : vector<4xf32> -> f32 - %inverse_sum0_f32 = scalar.divf %c1_f32, %sum0_scalar : f32 - %inverse_sum1_f32 = scalar.divf %c1_f32, %sum1_scalar : f32 - %inverse_sum2_f32 = scalar.divf %c1_f32, %sum2_scalar : f32 - %inverse_sum3_f32 = scalar.divf %c1_f32, %sum3_scalar : f32 - %inverse_sum0_f16 = scalar.fptrunc %inverse_sum0_f32 : f32 to f16 - %inverse_sum1_f16 = scalar.fptrunc %inverse_sum1_f32 : f32 to f16 - %inverse_sum2_f16 = scalar.fptrunc %inverse_sum2_f32 : f32 to f16 - %inverse_sum3_f16 = scalar.fptrunc %inverse_sum3_f32 : f32 to f16 - %inverse_sum0 = vector.splat %inverse_sum0_f16 : vector<4xf16> - %inverse_sum1 = vector.splat %inverse_sum1_f16 : vector<4xf16> - %inverse_sum2 = vector.splat %inverse_sum2_f16 : vector<4xf16> - %inverse_sum3 = vector.splat %inverse_sum3_f16 : vector<4xf16> - %normalized0_f16 = vector.mulf %final_output0, %inverse_sum0 : vector<4xf16> - %normalized1_f16 = vector.mulf %final_output1, %inverse_sum1 : vector<4xf16> - %normalized2_f16 = vector.mulf %final_output2, %inverse_sum2 : vector<4xf16> - %normalized3_f16 = vector.mulf %final_output3, %inverse_sum3 : vector<4xf16> - %normalized0 = vector.extf %normalized0_f16 : vector<4xf16> to vector<4xf32> - %normalized1 = vector.extf %normalized1_f16 : vector<4xf16> to vector<4xf32> - %normalized2 = vector.extf %normalized2_f16 : vector<4xf16> to vector<4xf32> - %normalized3 = vector.extf %normalized3_f16 : vector<4xf16> to vector<4xf32> - scf.if %query_valid0 { - vector.store %normalized0, %output_view[%query_head0, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_valid1 { - vector.store %normalized1, %output_view[%query_head1, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_valid2 { - vector.store %normalized2, %output_view[%query_head2, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_valid3 { - vector.store %normalized3, %output_view[%query_head3, %output_packet_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - } - kernel.return -} - -// The mask selects the first KV row exactly. QK, F16 probability conversion, -// and P*V all execute, while the expected result remains an auditable iota. -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row_case { - %key_value_token_count = check.literal value(64) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<64x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<64xf16> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<64xf16>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -// Thirty-two query heads share four KV heads in the production GQA ratio. -// Equal scores and constant values make every output exactly two while all -// head-group addressing and output ownership execute. -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x4x128xf16> - %value = check.generate.fill value(2.0) : tensor<128x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<128xf16> - %output = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(2.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -// Sixty-five KV rows force one masked cleanup tile after a full WMMA block. -// The mask selects only that final row, whose iota values begin at 1024. -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_tail_case { - %key_value_token_count = check.literal value(65) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> - %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<65xf16> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<65xf16>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -check.case public @qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case { - %key_value_token_count = check.param.choice values([64, 65, 128, 256, 512, 768, 1024, 1280, 2048]) name("key_value_token_count") : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_selected_row - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_gqa - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_tail_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_tail - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_64 {key_value_token_count = 64} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_65 {key_value_token_count = 65} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_128 {key_value_token_count = 128} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_256 {key_value_token_count = 256} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_512 {key_value_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_768 {key_value_token_count = 768} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_1024 {key_value_token_count = 1024} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_1280 {key_value_token_count = 1280} - -check.benchmark<@qwen3_moe_flash_attention_decode_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_f32_f16_wmma_decode_2048 {key_value_token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom deleted file mode 100644 index 5a0c8037f856..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom +++ /dev/null @@ -1,266 +0,0 @@ -// Exact-q128 grouped-query decode FlashAttention. -// -// One eight-wave workgroup owns each KV head. The waves cooperatively cover -// all 128 KV rows for QK, normalize the complete score matrix in LDS, and then -// each own one 16-channel P*V tile. This preserves 32 active waves for the -// production 32Q/4KV shape while removing duplicate QK work, global partial -// tensors, completion atomics, and a separate resolve dispatch. -amdgpu.target @qwen3_moe_attention_decode_q128_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] - -config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] - -kernel.def target(@qwen3_moe_attention_decode_q128_gfx11_wave64) @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %c1 = index.constant 1 : index - %c512 = index.constant 512 : index - kernel.launch.config workgroups(%key_value_head_count, %c1, %c1) workgroup_size(%c512, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) { - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_x0 = kernel.workgroup.id : index - %key_value_head = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 63)] : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 511)] : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane0 = kernel.subgroup.lane.id : index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : 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 - %c32 = index.constant 32 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %c512 = index.constant 512 : index - %c0_offset = index.constant 0 : offset - %query_stage_bytes = index.constant 4352 : offset - %score_stage_bytes = index.constant 12288 : offset - %probability_stage_bytes = index.constant 6144 : offset - %product_stage_bytes = index.constant 8192 : offset - %c0_f16 = scalar.constant 0.0 : f16 - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %c0_f32x2 = vector.constant 0.0 : vector<2xf32> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %head_size_f32 = scalar.constant 128.0 : f32 - %attention_scale = scalar.rsqrtf %head_size_f32 : f32 - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %key_value_width = index.mul %key_value_head_count, %c128 : index - %key_value_head_base = index.mul %key_value_head, %c128 : index - %score_key_origin = index.mul %subgroup, %c16 : index - %query_row_base = index.mul %subgroup, %c2 : index - %query_row0 = index.add %query_row_base, %c0 : index - %query_row1 = index.add %query_row_base, %c1 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head_base, %query_row1 : index - %query_row_valid0 = index.cmp ult, %query_row0, %query_heads_per_key_value_head : index - %query_row_valid1 = index.cmp ult, %query_row1, %query_heads_per_key_value_head : index - %query_head_in_range0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_in_range1 = index.cmp ult, %query_head1, %query_head_count : index - %query_valid0 = scalar.andi %query_row_valid0, %query_head_in_range0 : i1 - %query_valid1 = scalar.andi %query_row_valid1, %query_head_in_range1 : i1 - %query_valid0x2 = vector.from_elements %query_valid0, %query_valid0 : vector<2xi1> - %query_valid1x2 = vector.from_elements %query_valid1, %query_valid1 : vector<2xi1> - %key_token0 = index.add %lane, %c0 : index - %key_token1 = index.add %lane, %c64 : index - %output_tile_channel = index.mul %subgroup, %c16 : index - %query_transposed_layout = encoding.layout.strided [1, 136] : encoding - %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding - %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer - %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer - %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer - %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer - %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<128x[%key_value_width]xf16> - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<128x[%key_value_width]xf16> - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<128xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %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 - %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer - %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> - %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> - %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<128x24xf32> - %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x128xf16, %probability_transposed_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x128xf32> - // Scale and transpose the 16-row Q tile once for all eight score waves. - scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { - %linear = index.madd %load_iteration, %c512, %workitem : index - %local_query_row = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %local_query_head = index.add %query_head_base, %local_query_row : index - %local_query_row_valid = index.cmp ult, %local_query_row, %query_heads_per_key_value_head : index - %local_query_head_in_range = index.cmp ult, %local_query_head, %query_head_count : index - %local_query_valid = scalar.andi %local_query_row_valid, %local_query_head_in_range : i1 - %query_value = scf.if %local_query_valid -> (f16) { - %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 - %scaled = scalar.mulf %loaded, %attention_scale : f32 - %truncated = scalar.fptrunc %scaled : f32 to f16 - scf.yield %truncated : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Each wave produces one 16x16 score slice, covering all 128 KV rows. - %score_init_values = vector.constant 0.0 : vector<4xf32> - %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> - %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_channel = index.add %key_value_head_base, %head_tile : index - %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<128x[%key_value_width]xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %score_fragment, %score_stage_view[%score_key_origin, %c0] shape [%m, %n] : vector<4xf32>, view<128x24xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Each wave normalizes two query rows. Every lane owns two KV columns, so a - // vector reduction first combines its pair and the subgroup reduction then - // spans the complete 128-token row. - %mask0_f16 = view.load %mask_view[%key_token0] : view<128xf16> -> f16 - %mask1_f16 = view.load %mask_view[%key_token1] : view<128xf16> -> f16 - %mask0 = scalar.extf %mask0_f16 : f16 to f32 - %mask1 = scalar.extf %mask1_f16 : f16 to f32 - %mask_pair = vector.from_elements %mask0, %mask1 : vector<2xf32> - %raw_score00 = view.load %score_stage_view[%key_token0, %query_row0] : view<128x24xf32> -> f32 - %raw_score01 = view.load %score_stage_view[%key_token1, %query_row0] : view<128x24xf32> -> f32 - %raw_score10 = view.load %score_stage_view[%key_token0, %query_row1] : view<128x24xf32> -> f32 - %raw_score11 = view.load %score_stage_view[%key_token1, %query_row1] : view<128x24xf32> -> f32 - %raw_scores0 = vector.from_elements %raw_score00, %raw_score01 : vector<2xf32> - %raw_scores1 = vector.from_elements %raw_score10, %raw_score11 : vector<2xf32> - %added_scores0 = vector.addf %raw_scores0, %mask_pair : vector<2xf32> - %added_scores1 = vector.addf %raw_scores1, %mask_pair : vector<2xf32> - %negative_pair = vector.constant -1e+30 : vector<2xf32> - %masked_scores0 = vector.select %query_valid0x2, %added_scores0, %negative_pair : vector<2xf32> - %masked_scores1 = vector.select %query_valid1x2, %added_scores1, %negative_pair : vector<2xf32> - %lane_max0 = vector.reduce %masked_scores0, %negative_large : vector<2xf32>, f32 - %lane_max1 = vector.reduce %masked_scores1, %negative_large : vector<2xf32>, f32 - %lane_maxima = vector.from_elements %lane_max0, %lane_max1 : vector<2xf32> - %row_maxima = kernel.subgroup.reduce %lane_maxima : vector<2xf32> - %row_maximum0 = vector.extract %row_maxima[0] : vector<2xf32> -> f32 - %row_maximum1 = vector.extract %row_maxima[1] : vector<2xf32> -> f32 - %row_maximum0x2 = vector.splat %row_maximum0 : vector<2xf32> - %row_maximum1x2 = vector.splat %row_maximum1 : vector<2xf32> - %score_delta0 = vector.subf %masked_scores0, %row_maximum0x2 : vector<2xf32> - %score_delta1 = vector.subf %masked_scores1, %row_maximum1x2 : vector<2xf32> - %raw_probability0 = vector.expf %score_delta0 : vector<2xf32> - %raw_probability1 = vector.expf %score_delta1 : vector<2xf32> - %probability0 = vector.select %query_valid0x2, %raw_probability0, %c0_f32x2 : vector<2xf32> - %probability1 = vector.select %query_valid1x2, %raw_probability1, %c0_f32x2 : vector<2xf32> - %lane_sum0 = vector.reduce %probability0, %c0_f32 : vector<2xf32>, f32 - %lane_sum1 = vector.reduce %probability1, %c0_f32 : vector<2xf32>, f32 - %lane_sums = vector.from_elements %lane_sum0, %lane_sum1 : vector<2xf32> - %row_sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> - %row_sum0 = vector.extract %row_sums[0] : vector<2xf32> -> f32 - %row_sum1 = vector.extract %row_sums[1] : vector<2xf32> -> f32 - %reciprocal_sum0 = scf.if %query_valid0 -> (f32) { - %reciprocal = scalar.divf %c1_f32, %row_sum0 : f32 - scf.yield %reciprocal : f32 - } else { - scf.yield %c0_f32 : f32 - } - %reciprocal_sum1 = scf.if %query_valid1 -> (f32) { - %reciprocal = scalar.divf %c1_f32, %row_sum1 : f32 - scf.yield %reciprocal : f32 - } else { - scf.yield %c0_f32 : f32 - } - %reciprocal_sum0x2 = vector.splat %reciprocal_sum0 : vector<2xf32> - %reciprocal_sum1x2 = vector.splat %reciprocal_sum1 : vector<2xf32> - %normalized_probability0 = vector.mulf %probability0, %reciprocal_sum0x2 : vector<2xf32> - %normalized_probability1 = vector.mulf %probability1, %reciprocal_sum1x2 : vector<2xf32> - %probability0_f16 = vector.fptrunc %normalized_probability0 : vector<2xf32> to vector<2xf16> - %probability1_f16 = vector.fptrunc %normalized_probability1 : vector<2xf32> to vector<2xf16> - %probability00 = vector.extract %probability0_f16[0] : vector<2xf16> -> f16 - %probability01 = vector.extract %probability0_f16[1] : vector<2xf16> -> f16 - %probability10 = vector.extract %probability1_f16[0] : vector<2xf16> -> f16 - %probability11 = vector.extract %probability1_f16[1] : vector<2xf16> -> f16 - view.store %probability00, %probability_stage_view[%query_row0, %key_token0] : f16, view<16x128xf16, %probability_transposed_layout> - view.store %probability01, %probability_stage_view[%query_row0, %key_token1] : f16, view<16x128xf16, %probability_transposed_layout> - view.store %probability10, %probability_stage_view[%query_row1, %key_token0] : f16, view<16x128xf16, %probability_transposed_layout> - view.store %probability11, %probability_stage_view[%query_row1, %key_token1] : f16, view<16x128xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Each wave owns one disjoint 16-channel output tile. - %product_init_values = vector.constant 0.0 : vector<4xf32> - %product_init = vector.fragment %product_init_values shape [%m, %n] : vector<4xf32> - %value_channel = index.add %key_value_head_base, %output_tile_channel : index - %product_fragment = scf.for %key_tile = [%c0 to %c128 step %c16](%product_accumulator = %product_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x128xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%key_tile, %value_channel] shape [%k, %n] : view<128x[%key_value_width]xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_product_accumulator : vector<4xf32> - } - vector.fragment.store %product_fragment, %product_stage_view[%c0, %output_tile_channel] shape [%m, %n] : vector<4xf32>, view<16x128xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // One vector packet per workitem drains the 16x128 result tile while - // suppressing padded GQA rows. - %output_row = index.div %workitem, %c32 : index - %output_packet = index.rem %workitem, %c32 : index - %output_channel = index.mul %output_packet, %c4 : index - %output_head = index.add %query_head_base, %output_row : index - %output_row_valid = index.cmp ult, %output_row, %query_heads_per_key_value_head : index - %output_head_in_range = index.cmp ult, %output_head, %query_head_count : index - %output_valid = scalar.andi %output_row_valid, %output_head_in_range : i1 - scf.if %output_valid { - %output_packet_value = vector.load %product_stage_view[%output_row, %output_channel] : view<16x128xf32> -> vector<4xf32> - vector.store %output_packet_value, %output_view[%output_head, %output_channel] : vector<4xf32>, view<[%query_head_count]x128xf32> - } - kernel.return -} - -// A sharp mask selects the first KV row and makes the expected result an -// auditable iota while QK, global softmax, F16 probabilities, and P*V execute. -check.case public @qwen3_moe_flash_attention_decode_q128_fused_selected_row_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<128x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<128xf16> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<1x128xf32>, tensor<128x1x128xf16>, tensor<128x1x128xf16>, tensor<128xf16>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -// The production 32Q/4KV GQA shape verifies all head groups and output tiles. -check.case public @qwen3_moe_flash_attention_decode_q128_fused_gqa_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x4x128xf16> - %value = check.generate.fill value(2.0) : tensor<128x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<128xf16> - %output = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(2.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -check.case public @qwen3_moe_flash_attention_decode_q128_fused_benchmark_case { - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<128x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<128x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<128xf16> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_q128_fused_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %output) : [index](index, tensor<32x128xf32>, tensor<128x4x128xf16>, tensor<128x4x128xf16>, tensor<128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -check.benchmark<@qwen3_moe_flash_attention_decode_q128_fused_benchmark_case> @qwen3_moe_flash_attention_decode_q128_fused diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom deleted file mode 100644 index f07bdfb8ac53..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom +++ /dev/null @@ -1,1213 +0,0 @@ -// Qwen3 MoE grouped-query decode FlashAttention. -// -// Each workgroup processes one 64-token KV block for all GQA query heads that -// share a KV head. The workgroups publish online-softmax state, then the last -// arrival folds every block and resets the per-KV-head completion counter for -// the next invocation. Packing GQA heads removes redundant K/V traffic while -// split-K preserves enough parallelism for decode without another dispatch. -amdgpu.target @qwen3_moe_decode_split_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] - -config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] - -// Maximum K/V storage capacity available to the compiled kernel. -config.decl @qwen3_moe.attention.key_value_token_capacity : %value: index where [range(%value, 64, 32768)] - -// Computes one active online-softmax partial for a 64-row KV block. Both the -// fused short-context export and the two-dispatch long-context export reach -// this body through the block-classifying producer below, keeping their -// different binding contracts honest without duplicating the attention math. -// Partial storage retains its capacity-specialized block-axis stride while the -// launched block count bounds issue-time work. Keeping those values distinct -// preserves constant address arithmetic across changing visible prefixes. -func.template device @qwen3_moe_flash_attention_decode_split_produce_active_partials_body_f32_f16_wmma(%key_value_token_count: index, %partial_block_capacity0: index, %launched_block_count0: index, %query: buffer, %key: buffer, %value: buffer, %lane_mask: f32, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { - %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %partial_block_capacity, %launched_block_count = index.assume %partial_block_capacity0, %launched_block_count0 [range(%partial_block_capacity0, 1, 512), range(%launched_block_count0, 1, 512), le(%launched_block_count0, %partial_block_capacity0)] : index, index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_x0 = kernel.workgroup.id : index - %workgroup_y0 = kernel.workgroup.id : index - %workgroup_x_in_launch = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 511)] : index - %workgroup_y = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : 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 - %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 - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %query_stage_bytes = index.constant 4352 : offset - %score_stage_bytes = index.constant 6144 : offset - %probability_stage_bytes = index.constant 3072 : offset - %product_stage_bytes = index.constant 4096 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %head_size_f32 = scalar.constant 128.0 : f32 - %attention_scale = scalar.rsqrtf %head_size_f32 : f32 - %c0_f16 = scalar.constant 0.0 : f16 - %c0_f16x8 = vector.constant 0.0 : vector<8xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %workgroup_x, %launch_partial_block_capacity, %launch_launched_block_count = index.assume %workgroup_x_in_launch, %partial_block_capacity, %launched_block_count [lt(%workgroup_x_in_launch, %launched_block_count)] : index, index, index - %tail_key_value_token_count = index.rem %bounded_key_value_token_count, %c64 : index - %has_no_tail = index.cmp eq, %tail_key_value_token_count, %c0 : index - %active_padded_key_value_token_count = index.add %bounded_key_value_token_count, %c63 : index - %active_key_value_block_count = index.div %active_padded_key_value_token_count, %c64 : index - %last_block_ordinal = index.sub %active_key_value_block_count, %c1 : index - %block_ordinal = index.add %workgroup_x, %c0 : index - %is_not_last_block = index.cmp ne, %block_ordinal, %last_block_ordinal : index - %is_full_block = scalar.ori %has_no_tail, %is_not_last_block : i1 - %key_value_head = index.add %workgroup_y, %c0 : index - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %key_value_width = index.mul %key_value_head_count, %c128 : index - %key_value_head_base = index.mul %key_value_head, %c128 : index - %key_origin = index.mul %block_ordinal, %c64 : index - %subgroup_score_column = index.mul %subgroup, %c16 : index - %subgroup_query_row = index.mul %subgroup, %c4 : index - %query_row0 = index.add %subgroup_query_row, %c0 : index - %query_row1 = index.add %subgroup_query_row, %c1 : index - %query_row2 = index.add %subgroup_query_row, %c2 : index - %query_row3 = index.add %subgroup_query_row, %c3 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head_base, %query_row1 : index - %query_head2 = index.add %query_head_base, %query_row2 : index - %query_head3 = index.add %query_head_base, %query_row3 : index - %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index - %query_head_valid2 = index.cmp ult, %query_head2, %query_head_count : index - %query_head_valid3 = index.cmp ult, %query_head3, %query_head_count : index - %query_valid = vector.from_elements %query_head_valid0, %query_head_valid1, %query_head_valid2, %query_head_valid3 : vector<4xi1> - %subgroup_output_channel = index.mul %subgroup, %c32 : index - %subgroup_output_channel1 = index.add %subgroup_output_channel, %c16 : index - %lane_output_channel = index.mul %lane, %c4 : index - %lane_has_output = index.cmp ult, %lane, %c32 : index - %lane_is_zero = index.cmp eq, %lane, %c0 : index - %query_transposed_layout = encoding.layout.strided [1, 136] : encoding - %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding - %query_noalias, %key_noalias, %value_noalias, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias = buffer.assume.noalias %query, %key, %value, %partial_max, %partial_sum, %partial_output : buffer, buffer, buffer, buffer, buffer, buffer - %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer - %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer - %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer - %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer - %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer - %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer - %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> - %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 - %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer - %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> - %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> - %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> - %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x128xf16> - %tail_key_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<64x16xf16> - %tail_value_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<16x128xf16> - scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %local_query_row = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %local_query_head = index.add %query_head_base, %local_query_row : index - %local_query_head_valid = index.cmp ult, %local_query_head, %query_head_count : index - %query_value = scf.if %local_query_head_valid -> (f16) { - %loaded = view.load %query_view[%local_query_head, %query_channel] : view<[%query_head_count]x128xf32> -> f32 - %scaled = scalar.mulf %loaded, %attention_scale : f32 - %truncated = scalar.fptrunc %scaled : f32 to f16 - scf.yield %truncated : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %score_init_values = vector.constant 0.0 : vector<4xf32> - %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> - %score_fragment = scf.if %is_full_block -> (vector<4xf32>) { - %full_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - // Bound this view by the proven tile end so the full vector footprint is - // visible without treating physical tail padding as logical storage. - %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index - %score_key_end0 = index.add %score_key_origin0, %c16 : index - %score_key_origin, %score_key_end = index.assume %score_key_origin0, %score_key_end0 [le(%score_key_end0, %bounded_key_value_token_count)] : index, index - %full_key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%score_key_end]x[%key_value_width]xf16> - %key_channel = index.add %key_value_head_base, %head_tile : index - %key_fragment = vector.fragment.load %full_key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%score_key_end]x[%key_value_width]xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - scf.yield %full_score_fragment : vector<4xf32> - } else { - // Stage one 64x16 K panel at a time. Every physical load is guarded, so - // callers need no initialized padding beyond the logical KV length. - %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - scf.for %load_iteration = [%c0 to %c4 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %key_row = index.div %linear, %c16 : index - %head_channel = index.rem %linear, %c16 : index - %key_token = index.add %key_origin, %key_row : index - %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index - %key_value = scf.if %key_valid -> (f16) { - %head_channel0 = index.add %head_tile, %head_channel : index - %key_channel = index.add %key_value_head_base, %head_channel0 : index - %loaded = view.load %key_view[%key_token, %key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %key_value, %tail_key_stage_view[%key_row, %head_channel] : f16, view<64x16xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %key_fragment = vector.fragment.load %tail_key_stage_view[%subgroup_score_column, %c0] shape [%m, %k] : view<64x16xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next_score_accumulator : vector<4xf32> - } - scf.yield %tail_score_fragment : vector<4xf32> - } - vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %key_token = index.add %key_origin, %lane : index - %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 - %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 - %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 - %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 - %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index - %score_valid0 = scalar.andi %query_head_valid0, %key_valid : i1 - %score_valid1 = scalar.andi %query_head_valid1, %key_valid : i1 - %score_valid2 = scalar.andi %query_head_valid2, %key_valid : i1 - %score_valid3 = scalar.andi %query_head_valid3, %key_valid : i1 - %masked_score0 = scf.if %score_valid0 -> (f32) { - %score = scalar.addf %raw_score0, %lane_mask : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score1 = scf.if %score_valid1 -> (f32) { - %score = scalar.addf %raw_score1, %lane_mask : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score2 = scf.if %score_valid2 -> (f32) { - %score = scalar.addf %raw_score2, %lane_mask : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_score3 = scf.if %score_valid3 -> (f32) { - %score = scalar.addf %raw_score3, %lane_mask : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %score_delta = vector.subf %masked_scores, %block_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Keep both P*V halves live so the scheduler can interleave their independent - // matrix chains. F16 accumulation and exchange preserve the Vulkan CM1 - // contract while halving the product LDS footprint. - %product_init0 = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %product_init1 = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %value_channel0 = index.add %key_value_head_base, %subgroup_output_channel : index - %value_channel1 = index.add %key_value_head_base, %subgroup_output_channel1 : index - %product_fragment0, %product_fragment1 = scf.if %is_full_block -> (vector<8xf16>, vector<8xf16>) { - %full_product_fragment0, %full_product_fragment1 = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator0 = %product_init0 : vector<8xf16>, %product_accumulator1 = %product_init1 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>) unroll { - // Retain the same logical-extent contract for the full P*V tile. - %value_token0 = index.add %key_origin, %key_tile : index - %value_token_end0 = index.add %value_token0, %c16 : index - %value_token, %value_token_end = index.assume %value_token0, %value_token_end0 [le(%value_token_end0, %bounded_key_value_token_count)] : index, index - %full_value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%value_token_end]x[%key_value_width]xf16> - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment0 = vector.fragment.load %full_value_view[%value_token, %value_channel0] shape [%k, %n] : view<[%value_token_end]x[%key_value_width]xf16> -> vector<16xf16> - %value_fragment1 = vector.fragment.load %full_value_view[%value_token, %value_channel1] shape [%k, %n] : view<[%value_token_end]x[%key_value_width]xf16> -> vector<16xf16> - %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> - } - scf.yield %full_product_fragment0, %full_product_fragment1 : vector<8xf16>, vector<8xf16> - } else { - // Stage one 16x128 V panel at a time and zero every lane beyond the - // logical KV length before the panel participates in P*V. - %tail_product_fragment0, %tail_product_fragment1 = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator0 = %product_init0 : vector<8xf16>, %product_accumulator1 = %product_init1 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>) unroll { - scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %key_row = index.div %linear, %c128 : index - %value_channel = index.rem %linear, %c128 : index - %value_token0 = index.add %key_origin, %key_tile : index - %value_token = index.add %value_token0, %key_row : index - %value_valid = index.cmp ult, %value_token, %bounded_key_value_token_count : index - %value_element = scf.if %value_valid -> (f16) { - %global_value_channel = index.add %key_value_head_base, %value_channel : index - %loaded = view.load %value_view[%value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %value_element, %tail_value_stage_view[%key_row, %value_channel] : f16, view<16x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment0 = vector.fragment.load %tail_value_stage_view[%c0, %subgroup_output_channel] shape [%k, %n] : view<16x128xf16> -> vector<16xf16> - %value_fragment1 = vector.fragment.load %tail_value_stage_view[%c0, %subgroup_output_channel1] shape [%k, %n] : view<16x128xf16> -> vector<16xf16> - %next_product_accumulator0 = vector.mma %probability_fragment, %value_fragment0, %product_accumulator0 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %next_product_accumulator1 = vector.mma %probability_fragment, %value_fragment1, %product_accumulator1 : vector<16xf16>, vector<16xf16>, vector<8xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next_product_accumulator0, %next_product_accumulator1 : vector<8xf16>, vector<8xf16> - } - scf.yield %tail_product_fragment0, %tail_product_fragment1 : vector<8xf16>, vector<8xf16> - } - vector.fragment.store %product_fragment0, %product_stage_view[%c0, %subgroup_output_channel] shape [%m, %n] : vector<8xf16>, view<16x128xf16> - vector.fragment.store %product_fragment1, %product_stage_view[%c0, %subgroup_output_channel1] shape [%m, %n] : vector<8xf16>, view<16x128xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.if %lane_has_output { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_output_channel] : view<16x128xf16> -> vector<4xf16> - scf.if %query_head_valid0 { - vector.store %block_output0, %partial_output_view[%key_value_head, %block_ordinal, %query_row0, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> - } - scf.if %query_head_valid1 { - vector.store %block_output1, %partial_output_view[%key_value_head, %block_ordinal, %query_row1, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> - } - scf.if %query_head_valid2 { - vector.store %block_output2, %partial_output_view[%key_value_head, %block_ordinal, %query_row2, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> - } - scf.if %query_head_valid3 { - vector.store %block_output3, %partial_output_view[%key_value_head, %block_ordinal, %query_row3, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16x128xf16> - } - } - scf.if %lane_is_zero { - scf.if %query_head_valid0 { - %maximum = vector.extract %block_max[0] : vector<4xf32> -> f32 - %sum = vector.extract %block_sum[0] : vector<4xf32> -> f32 - view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - } - scf.if %query_head_valid1 { - %maximum = vector.extract %block_max[1] : vector<4xf32> -> f32 - %sum = vector.extract %block_sum[1] : vector<4xf32> -> f32 - view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - } - scf.if %query_head_valid2 { - %maximum = vector.extract %block_max[2] : vector<4xf32> -> f32 - %sum = vector.extract %block_sum[2] : vector<4xf32> -> f32 - view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - } - scf.if %query_head_valid3 { - %maximum = vector.extract %block_max[3] : vector<4xf32> -> f32 - %sum = vector.extract %block_sum[3] : vector<4xf32> -> f32 - view.store %maximum, %partial_max_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - view.store %sum, %partial_sum_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%launch_partial_block_capacity]x16xf32> - } - } - func.return -} - -// Fully masked splits publish the online-softmax identity without running QK -// or P*V. Valid additive masks contain finite F16 values or negative infinity, -// and every finite F16 value is greater than -1e30. Comparing after extension -// therefore distinguishes active rows from masked rows exactly. -func.template device @qwen3_moe_flash_attention_decode_split_produce_partials_body_f32_f16_wmma(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { - %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_x0 = kernel.workgroup.id : index - %workgroup_y0 = kernel.workgroup.id : index - %workgroup_x_in_launch = index.assume %workgroup_x0 [range(%workgroup_x0, 0, 511)] : index - %workgroup_y = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : 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 - %c2 = index.constant 2 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c32 = index.constant 32 : index - %c64 = index.constant 64 : index - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c63 = index.constant 63 : index - %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index - %key_value_block_count0 = index.div %padded_key_value_token_capacity, %c64 : index - %key_value_block_count = index.assume %key_value_block_count0 [range(%key_value_block_count0, 1, 512)] : index - %workgroup_x, %launch_key_value_block_count = index.assume %workgroup_x_in_launch, %key_value_block_count [lt(%workgroup_x_in_launch, %key_value_block_count)] : index, index - %block_ordinal = index.add %workgroup_x, %c0 : index - %key_origin = index.mul %block_ordinal, %c64 : index - %key_token = index.add %key_origin, %lane : index - %key_valid = index.cmp ult, %key_token, %bounded_key_value_token_count : index - %mask_aligned = buffer.assume.alignment %mask {minimum_alignment = 16} : buffer - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]xf16> - %lane_mask = scf.if %key_valid -> (f32) { - %mask_f16 = view.load %mask_view[%key_token] : view<[%bounded_key_value_token_count]xf16> -> f16 - %mask_f32 = scalar.extf %mask_f16 : f16 to f32 - scf.yield %mask_f32 : f32 - } else { - scf.yield %negative_large : f32 - } - %block_mask_maximum = kernel.workgroup.reduce %lane_mask : f32 - %block_has_attention = scalar.cmpf ogt, %block_mask_maximum, %negative_large : f32 - scf.if %block_has_attention { - func.call inline @qwen3_moe_flash_attention_decode_split_produce_active_partials_body_f32_f16_wmma(%bounded_key_value_token_count, %launch_key_value_block_count, %launch_key_value_block_count, %query, %key, %value, %lane_mask, %partial_max, %partial_sum, %partial_output) : (index, index, index, buffer, buffer, buffer, f32, buffer, buffer, buffer) - } else { - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %key_value_head = index.add %workgroup_y, %c0 : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %subgroup_query_row = index.mul %subgroup, %c4 : index - %query_row0 = index.add %subgroup_query_row, %c0 : index - %query_row1 = index.add %subgroup_query_row, %c1 : index - %query_row2 = index.add %subgroup_query_row, %c2 : index - %query_row3 = index.add %subgroup_query_row, %c3 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head_base, %query_row1 : index - %query_head2 = index.add %query_head_base, %query_row2 : index - %query_head3 = index.add %query_head_base, %query_row3 : index - %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index - %query_head_valid2 = index.cmp ult, %query_head2, %query_head_count : index - %query_head_valid3 = index.cmp ult, %query_head3, %query_head_count : index - %lane_output_channel = index.mul %lane, %c4 : index - %lane_has_output = index.cmp ult, %lane, %c32 : index - %lane_is_zero = index.cmp eq, %lane, %c0 : index - %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output : buffer, buffer, buffer - %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer - %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer - %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer - %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> - scf.if %lane_has_output { - scf.if %query_head_valid0 { - vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row0, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> - } - scf.if %query_head_valid1 { - vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row1, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> - } - scf.if %query_head_valid2 { - vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row2, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> - } - scf.if %query_head_valid3 { - vector.store %c0_f16x4, %partial_output_view[%key_value_head, %block_ordinal, %query_row3, %lane_output_channel] : vector<4xf16>, view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> - } - } - scf.if %lane_is_zero { - scf.if %query_head_valid0 { - view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row0] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - } - scf.if %query_head_valid1 { - view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row1] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - } - scf.if %query_head_valid2 { - view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row2] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - } - scf.if %query_head_valid3 { - view.store %negative_large, %partial_max_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - view.store %c0_f32, %partial_sum_view[%key_value_head, %block_ordinal, %query_row3] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - } - } - } - func.return -} - -// Up to four split-K blocks are cheapest to fold directly. Each workitem owns -// one output element, so the reducer needs no LDS or subgroup synchronization. -func.template device @qwen3_moe_flash_attention_decode_split_reduce_completed_direct_f32(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) where [range(%partial_block_capacity0, 1, 4)] { - %partial_block_capacity, %active_block_count = index.assume %partial_block_capacity0, %active_block_count0 [range(%partial_block_capacity0, 1, 4), range(%active_block_count0, 1, 4), le(%active_block_count0, %partial_block_capacity0)] : index, index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_y0 = kernel.workgroup.id : index - %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %partial_max_aligned = buffer.assume.alignment %partial_max {minimum_alignment = 16} : buffer - %partial_sum_aligned = buffer.assume.alignment %partial_sum {minimum_alignment = 16} : buffer - %partial_output_aligned = buffer.assume.alignment %partial_output {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output {minimum_alignment = 16} : buffer - %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> - %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> - %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %partial_element_count = index.mul %query_heads_per_key_value_head, %c128 : index - scf.for %linear = [%workitem to %partial_element_count step %c256] { - %query_row = index.div %linear, %c128 : index - %output_channel = index.rem %linear, %c128 : index - %query_head = index.add %query_head_base, %query_row : index - %query_head_valid = index.cmp ult, %query_head, %query_head_count : index - scf.if %query_head_valid { - %maximum = scf.for %block = [%c0 to %active_block_count step %c1](%running_maximum = %negative_large : f32) -> (f32) unroll schedule(interleaved) { - %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %next_maximum = scalar.maxnumf %running_maximum, %block_maximum : f32 - scf.yield %next_maximum : f32 - } - %sum, %unnormalized_output = scf.for %block = [%c0 to %active_block_count step %c1](%running_sum = %c0_f32 : f32, %running_output = %c0_f32 : f32) -> (f32, f32) unroll schedule(interleaved) { - %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %delta = scalar.subf %block_maximum, %maximum : f32 - %scale = scalar.expf %delta : f32 - %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %scaled_sum = scalar.mulf %block_sum, %scale : f32 - %next_sum = scalar.addf %running_sum, %scaled_sum : f32 - %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> -> f16 - %block_output = scalar.extf %block_output_f16 : f16 to f32 - %scaled_output = scalar.mulf %block_output, %scale : f32 - %next_output = scalar.addf %running_output, %scaled_output : f32 - scf.yield %next_sum, %next_output : f32, f32 - } - %normalized_output = scalar.divf %unnormalized_output, %sum : f32 - view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x128xf32> - } - } - func.return -} - -// Longer bounded contexts amortize a cooperative reducer. Each wave folds two -// query rows, stages the per-block scales in LDS, and writes two output channels -// per lane. -func.template device @qwen3_moe_flash_attention_decode_split_reduce_completed_cooperative_f32(%partial_block_capacity0: index, %active_block_count0: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) { - %partial_block_capacity, %active_block_count = index.assume %partial_block_capacity0, %active_block_count0 [range(%partial_block_capacity0, 1, 32), range(%active_block_count0, 1, 32), le(%active_block_count0, %partial_block_capacity0)] : index, index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_y0 = kernel.workgroup.id : index - %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : 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 - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c7 = index.constant 7 : index - %c8 = index.constant 8 : index - %c0_offset = index.constant 0 : offset - %scale_stage_bytes = index.constant 1024 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %c0_f32x2 = vector.constant 0.0 : vector<2xf32> - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %lane_output_channel0 = index.mul %lane, %c2 : index - %lane_output_channel = index.assume %lane_output_channel0 [range(%lane_output_channel0, 0, 126)] : index - %lane_has_block = index.cmp ult, %lane, %active_block_count : index - %partial_max_aligned = buffer.assume.alignment %partial_max {minimum_alignment = 16} : buffer - %partial_sum_aligned = buffer.assume.alignment %partial_sum {minimum_alignment = 16} : buffer - %partial_output_aligned = buffer.assume.alignment %partial_output {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output {minimum_alignment = 16} : buffer - %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> - %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> - %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer - %scale_stage_view = buffer.view %scale_stage[%c0_offset] : buffer -> view<4x32x2xf32> - %padded_reduction_row_count = index.add %query_heads_per_key_value_head, %c7 : index - %reduction_phase_count = index.div %padded_reduction_row_count, %c8 : index - scf.for %phase = [%c0 to %reduction_phase_count step %c1] unroll { - %phase_row_base = index.mul %phase, %c8 : index - %subgroup_row_base = index.mul %subgroup, %c2 : index - %query_row0 = index.add %phase_row_base, %subgroup_row_base : index - %query_row1 = index.add %query_row0, %c1 : index - %query_head0 = index.add %query_head_base, %query_row0 : index - %query_head1 = index.add %query_head0, %c1 : index - %query_head_valid0 = index.cmp ult, %query_head0, %query_head_count : index - %query_head_valid1 = index.cmp ult, %query_head1, %query_head_count : index - %safe_query_row0 = scf.select %query_head_valid0, %query_row0, %c0 : index - %safe_query_row1 = scf.select %query_head_valid1, %query_row1, %c0 : index - %reducer_active0 = scalar.andi %lane_has_block, %query_head_valid0 : i1 - %reducer_active1 = scalar.andi %lane_has_block, %query_head_valid1 : i1 - %local_maximum0 = scf.if %reducer_active0 -> (f32) { - %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - scf.yield %block_maximum : f32 - } else { - scf.yield %negative_large : f32 - } - %local_maximum1 = scf.if %reducer_active1 -> (f32) { - %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - scf.yield %block_maximum : f32 - } else { - scf.yield %negative_large : f32 - } - %local_maximums = vector.from_elements %local_maximum0, %local_maximum1 : vector<2xf32> - %maximums = kernel.subgroup.reduce %local_maximums : vector<2xf32> - %maximum0 = vector.extract %maximums[0] : vector<2xf32> -> f32 - %maximum1 = vector.extract %maximums[1] : vector<2xf32> -> f32 - %local_sum0, %scale0 = scf.if %reducer_active0 -> (f32, f32) { - %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %delta = scalar.subf %block_maximum, %maximum0 : f32 - %scale = scalar.expf %delta : f32 - %block_sum = view.load %partial_sum_view[%key_value_head, %lane, %query_row0] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %scaled_sum = scalar.mulf %block_sum, %scale : f32 - scf.yield %scaled_sum, %scale : f32, f32 - } else { - scf.yield %c0_f32, %c0_f32 : f32, f32 - } - %local_sum1, %scale1 = scf.if %reducer_active1 -> (f32, f32) { - %block_maximum = view.load %partial_max_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %delta = scalar.subf %block_maximum, %maximum1 : f32 - %scale = scalar.expf %delta : f32 - %block_sum = view.load %partial_sum_view[%key_value_head, %lane, %query_row1] : view<[%key_value_head_count]x[%partial_block_capacity]x16xf32> -> f32 - %scaled_sum = scalar.mulf %block_sum, %scale : f32 - scf.yield %scaled_sum, %scale : f32, f32 - } else { - scf.yield %c0_f32, %c0_f32 : f32, f32 - } - scf.if %lane_has_block { - %bounded_block_lane = index.assume %lane [range(%lane, 0, 31)] : index - %scales = vector.from_elements %scale0, %scale1 : vector<2xf32> - vector.store %scales, %scale_stage_view[%subgroup, %bounded_block_lane, %c0] : vector<2xf32>, view<4x32x2xf32> - } - %local_sums = vector.from_elements %local_sum0, %local_sum1 : vector<2xf32> - %sums = kernel.subgroup.reduce %local_sums : vector<2xf32> - %sum0 = vector.extract %sums[0] : vector<2xf32> -> f32 - %sum1 = vector.extract %sums[1] : vector<2xf32> -> f32 - kernel.barrier scope(workgroup) ordering(acq_rel) - %unnormalized_output0, %unnormalized_output1 = scf.for %block = [%c0 to %active_block_count step %c1](%running_output0 = %c0_f32x2 : vector<2xf32>, %running_output1 = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>, vector<2xf32>) unroll(%c4) schedule(interleaved) { - %scales = vector.load %scale_stage_view[%subgroup, %block, %c0] : view<4x32x2xf32> -> vector<2xf32> - %scale0 = vector.extract %scales[0] : vector<2xf32> -> f32 - %scale1 = vector.extract %scales[1] : vector<2xf32> -> f32 - %block_output0_f16 = vector.load %partial_output_view[%key_value_head, %block, %safe_query_row0, %lane_output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> -> vector<2xf16> - %block_output0 = vector.extf %block_output0_f16 : vector<2xf16> to vector<2xf32> - %scale_vector0 = vector.splat %scale0 : vector<2xf32> - %scaled_output0 = vector.mulf %block_output0, %scale_vector0 : vector<2xf32> - %next_output0 = vector.addf %running_output0, %scaled_output0 : vector<2xf32> - %block_output1_f16 = vector.load %partial_output_view[%key_value_head, %block, %safe_query_row1, %lane_output_channel] : view<[%key_value_head_count]x[%partial_block_capacity]x16x128xf16> -> vector<2xf16> - %block_output1 = vector.extf %block_output1_f16 : vector<2xf16> to vector<2xf32> - %scale_vector1 = vector.splat %scale1 : vector<2xf32> - %scaled_output1 = vector.mulf %block_output1, %scale_vector1 : vector<2xf32> - %next_output1 = vector.addf %running_output1, %scaled_output1 : vector<2xf32> - scf.yield %next_output0, %next_output1 : vector<2xf32>, vector<2xf32> - } - %sum_vector0 = vector.splat %sum0 : vector<2xf32> - %sum_vector1 = vector.splat %sum1 : vector<2xf32> - %normalized_output0 = vector.divf %unnormalized_output0, %sum_vector0 : vector<2xf32> - %normalized_output1 = vector.divf %unnormalized_output1, %sum_vector1 : vector<2xf32> - scf.if %query_head_valid0 { - vector.store %normalized_output0, %output_view[%query_head0, %lane_output_channel] : vector<2xf32>, view<[%query_head_count]x128xf32> - } - scf.if %query_head_valid1 { - vector.store %normalized_output1, %output_view[%query_head1, %lane_output_channel] : vector<2xf32>, view<[%query_head_count]x128xf32> - } - } - func.return -} - -// Packs the query heads owned by one completed KV head into GGML's Q8_1 x4 -// layout. Each 32-workitem cohort owns one contiguous 128-element query row; -// the production 8:1 GQA ratio therefore fills all 256 workitems. Smaller or -// larger valid ratios use the same phase loop without making barriers -// conditional. -func.template device @qwen3_moe_flash_attention_decode_pack_completed_key_value_head_q8_1_x4(%key_value_head: index, %output: buffer, %q8_output: buffer) { - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index - %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 - %c32 = index.constant 32 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 1024 : offset - %scratch_bytes = index.constant 1152 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %output_noalias, %q8_output_noalias = buffer.assume.noalias %output, %q8_output : buffer, buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %q8_output_aligned = buffer.assume.alignment %q8_output_noalias {minimum_alignment = 16} : buffer - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %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> - %group_in_phase = index.div %workitem, %c32 : index - %lane_in_group0 = index.rem %workitem, %c32 : index - %lane_in_group = index.assume %lane_in_group0 [range(%lane_in_group0, 0, 31)] : index - %block_in_group0 = index.div %lane_in_group, %c8 : index - %block_in_group = index.assume %block_in_group0 [range(%block_in_group0, 0, 3)] : index - %word_in_block0 = index.rem %lane_in_group, %c8 : index - %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 7)] : index - %padded_phase_count = index.add %query_heads_per_key_value_head, %c7 : index - %phase_count = index.div %padded_phase_count, %c8 : index - scf.for %phase = [%c0 to %phase_count step %c1] unroll { - %phase_group_base = index.mul %phase, %c8 : index - %group_in_key_value_head = index.add %phase_group_base, %group_in_phase : index - %valid_group = index.cmp ult, %group_in_key_value_head, %query_heads_per_key_value_head : index - %key_value_query_head_base = index.mul %key_value_head, %query_heads_per_key_value_head : index - %query_head0 = index.add %key_value_query_head_base, %group_in_key_value_head : index - %safe_query_head = scf.select %valid_group, %query_head0, %c0 : index - %block_element_add = index.mul %block_in_group, %c32 : index - %word_element_add = index.mul %word_in_block, %c4 : index - %input_channel0 = index.add %block_element_add, %word_element_add : index - %input_channel = index.assume %input_channel0 [range(%input_channel0, 0, 124)] : index - %input_values = scf.if %valid_group -> (vector<4xf32>) { - %values = vector.load %output_view[%safe_query_head, %input_channel] : view<[%query_head_count]x128xf32> -> vector<4xf32> - scf.yield %values : vector<4xf32> - } else { - scf.yield %c0_f32x4 : vector<4xf32> - } - %absolute_values = vector.absf %input_values : vector<4xf32> - %thread_max = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_max, %scratch_values[%workitem] : f32, view<256xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_cohort_leader = index.cmp eq, %word_in_block, %c0 : index - %d_index0 = index.mul %group_in_phase, %c4 : index - %d_index1 = index.add %d_index0, %block_in_group : index - %d_index = index.assume %d_index1 [range(%d_index1, 0, 31)] : index - scf.if %is_cohort_leader { - %group_thread_base = index.mul %group_in_phase, %c32 : index - %block_thread_add = index.mul %block_in_group, %c8 : index - %cohort_base = index.add %group_thread_base, %block_thread_add : index - %cohort_maxima = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> - %amax = vector.reduce %cohort_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%d_index] : f32, view<32xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %d = view.load %scratch_d[%d_index] : view<32xf32> -> f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %group_byte_offset = index.scale %safe_query_head, %group_bytes : index, offset -> offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %q8_output_aligned[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %q8_output_aligned[%payload_byte_offset] : buffer -> view<32xi32> - %block_word_add = index.mul %block_in_group, %c8 : index - %packed_word_index0 = index.add %block_word_add, %word_in_block : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - scf.if %valid_group { - vector.store %packed_word, %group_qs[%packed_word_index] : vector<1xi32>, view<32xi32> - } - %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %thread_sum, %scratch_values[%workitem] : f32, view<256xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.if %is_cohort_leader { - %group_thread_base = index.mul %group_in_phase, %c32 : index - %block_thread_add = index.mul %block_in_group, %c8 : index - %cohort_base = index.add %group_thread_base, %block_thread_add : index - %cohort_sums = vector.load %scratch_values[%cohort_base] : view<256xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %cohort_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - %s_index = index.add %ds_index, %c1 : index - scf.if %valid_group { - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - } - kernel.barrier scope(workgroup) ordering(acq_rel) - } - func.return -} - -// Completes bounded decode contexts inside the last arriving producer -// workgroup. Four KV-head workgroups reduce their own query heads while all -// other producers retire, erasing a second dispatch and its execution barrier. -// Capacity selects the algorithm and partial layout; producer count owns the -// issue-time completion threshold and active reduction prefix. -func.template device priority(20) @qwen3_moe_flash_attention_decode_split_reduce_fused_direct_f32(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) where [range(%key_value_token_capacity, 64, 256)] { - %partial_block_capacity, %producer_block_count = index.assume %partial_block_capacity0, %producer_block_count0 [range(%partial_block_capacity0, 1, 4), range(%producer_block_count0, 1, 4), le(%producer_block_count0, %partial_block_capacity0)] : index, index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_y0 = kernel.workgroup.id : index - %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %completion_counter_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %completion_counter, %output : buffer, buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - %workitem_is_zero = index.cmp eq, %workitem, %c0 : index - // Publish every producer's partial stores before the leader advances one - // workgroup arrival. The last arrival then acquires every partial. - kernel.barrier scope(workgroup) ordering(release) - scf.if %workitem_is_zero { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%key_value_head] {ordering = acq_rel, scope = device} : i32, view<[%key_value_head_count]xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %key_value_block_count_i32 = index.cast %producer_block_count : index to i32 - %last_block_ordinal_i32 = scalar.subi %key_value_block_count_i32, %c1_i32 : i32 - %negative_key_value_block_count_i32 = scalar.subi %c0_i32, %key_value_block_count_i32 : i32 - %is_last_partition = scalar.cmpi eq, %old_counter, %last_block_ordinal_i32 : i32 - scf.if %is_last_partition { - kernel.barrier scope(workgroup) ordering(acquire) - func.call inline @qwen3_moe_flash_attention_decode_split_reduce_completed_direct_f32(%partial_block_capacity, %producer_block_count, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias) : (index, index, buffer, buffer, buffer, buffer) - scf.if %publish_q8 { - kernel.barrier scope(workgroup) ordering(acq_rel) - func.apply(%key_value_head, %output_noalias, %q8_output) : (index, buffer, buffer) - } - // Do not expose the reset until every final F32 and Q8 store completes. - kernel.barrier scope(workgroup) ordering(release) - scf.if %workitem_is_zero { - view.atomic.reduce %negative_key_value_block_count_i32, %completion_counter_view[%key_value_head] {ordering = release, scope = device} : i32, view<[%key_value_head_count]xi32> - } - } - func.return -} - -// Contexts without a proven short bound use the cooperative completion path. -func.template device priority(10) @qwen3_moe_flash_attention_decode_split_reduce_fused_cooperative_f32(%key_value_token_capacity: index, %partial_block_capacity0: index, %producer_block_count0: index, %publish_q8: i1, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %q8_output: buffer) where [range(%key_value_token_capacity, 257, 2048)] { - %partial_block_capacity, %producer_block_count = index.assume %partial_block_capacity0, %producer_block_count0 [range(%partial_block_capacity0, 1, 32), range(%producer_block_count0, 1, 32), le(%producer_block_count0, %partial_block_capacity0)] : index, index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %workgroup_y0 = kernel.workgroup.id : index - %key_value_head = index.assume %workgroup_y0 [range(%workgroup_y0, 0, 63)] : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %completion_counter_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %completion_counter, %output : buffer, buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - %workitem_is_zero = index.cmp eq, %workitem, %c0 : index - // Publish every producer's partial stores before the leader advances one - // workgroup arrival. The last arrival then acquires every partial. - kernel.barrier scope(workgroup) ordering(release) - scf.if %workitem_is_zero { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%key_value_head] {ordering = acq_rel, scope = device} : i32, view<[%key_value_head_count]xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %key_value_block_count_i32 = index.cast %producer_block_count : index to i32 - %last_block_ordinal_i32 = scalar.subi %key_value_block_count_i32, %c1_i32 : i32 - %negative_key_value_block_count_i32 = scalar.subi %c0_i32, %key_value_block_count_i32 : i32 - %is_last_partition = scalar.cmpi eq, %old_counter, %last_block_ordinal_i32 : i32 - scf.if %is_last_partition { - kernel.barrier scope(workgroup) ordering(acquire) - func.call inline @qwen3_moe_flash_attention_decode_split_reduce_completed_cooperative_f32(%partial_block_capacity, %producer_block_count, %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias) : (index, index, buffer, buffer, buffer, buffer) - scf.if %publish_q8 { - kernel.barrier scope(workgroup) ordering(acq_rel) - func.apply(%key_value_head, %output_noalias, %q8_output) : (index, buffer, buffer) - } - // Do not expose the reset until every final F32 and Q8 store completes. - kernel.barrier scope(workgroup) ordering(release) - scf.if %workitem_is_zero { - view.atomic.reduce %negative_key_value_block_count_i32, %completion_counter_view[%key_value_head] {ordering = release, scope = device} : i32, view<[%key_value_head_count]xi32> - } - } - func.return -} - -// Short-context export: produce and reduce in one dispatch. -kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index - %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index - kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer) { - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 2048)] : index - %bounded_key_value_token_count, %launch_key_value_token_capacity = index.assume %key_value_token_count_in_range, %key_value_token_capacity [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index, index - %padded_key_value_token_capacity = index.add %launch_key_value_token_capacity, %c63 : index - %producer_block_count = index.div %padded_key_value_token_capacity, %c64 : index - %publish_q8 = scalar.constant false : i1 - func.apply(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - func.apply(%launch_key_value_token_capacity, %producer_block_count, %producer_block_count, %publish_q8, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %partial_output) : (index, index, index, i1, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Short-context export that publishes the F32 attention result and the Q8_1 -// representation consumed by the following output projection. -kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index - %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index - kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) { - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 2048)] : index - %bounded_key_value_token_count, %launch_key_value_token_capacity = index.assume %key_value_token_count_in_range, %key_value_token_capacity [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index, index - %padded_key_value_token_capacity = index.add %launch_key_value_token_capacity, %c63 : index - %producer_block_count = index.div %padded_key_value_token_capacity, %c64 : index - %publish_q8 = scalar.constant true : i1 - func.apply(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - func.apply(%launch_key_value_token_capacity, %producer_block_count, %producer_block_count, %publish_q8, %partial_max, %partial_sum, %partial_output, %completion_counter, %output, %next_q8_output) : (index, index, index, i1, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Long-context producer export: publish partials for a following parallel -// reducer without carrying short-context synchronization bindings. -kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma(%key_value_token_count: index) { - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index - %key_value_block_count = index.div %padded_key_value_token_capacity, %c64 : index - kernel.launch.config workgroups(%key_value_block_count, %key_value_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer) { - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %key_value_token_count_in_range = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %bounded_key_value_token_count = index.assume %key_value_token_count_in_range [le(%key_value_token_count_in_range, %key_value_token_capacity)] : index - func.apply(%bounded_key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : (index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Long contexts have enough split-K partials that assigning the reduction to -// one last-arriving producer serializes useful work. This reducer launches one -// two-wave workgroup per query head after an execution barrier from the -// producer. The first wave computes normalization once and both waves consume -// it, avoiding the duplicate work of independent 64-channel output slices. -kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_reduce_f32(%key_value_token_count: index) { - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %c1 = index.constant 1 : index - %c128 = index.constant 128 : index - kernel.launch.config workgroups(%c1, %query_head_count, %c1) workgroup_size(%c128, %c1, %c1) : index -} launch(%key_value_token_count: index, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %output: buffer) { - %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %query_head0 = kernel.workgroup.id : index - %query_head = index.assume %query_head0 [range(%query_head0, 0, 63)] : index - %workitem = kernel.workitem.id : index - %subgroup = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c64 = index.constant 64 : index - %c0_offset = index.constant 0 : offset - %reduction_stage_bytes = index.constant 8 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %key_value_token_capacity = config.get @qwen3_moe.attention.key_value_token_capacity : index - %c63 = index.constant 63 : index - %padded_key_value_token_capacity = index.add %key_value_token_capacity, %c63 : index - %key_value_block_count0 = index.div %padded_key_value_token_capacity, %c64 : index - %key_value_block_count = index.assume %key_value_block_count0 [range(%key_value_block_count0, 1, 512)] : index - %query_heads_per_key_value_head0 = index.div %query_head_count, %key_value_head_count : index - %query_heads_per_key_value_head = index.assume %query_heads_per_key_value_head0 [range(%query_heads_per_key_value_head0, 1, 16)] : index - %key_value_head = index.div %query_head, %query_heads_per_key_value_head : index - %query_row = index.rem %query_head, %query_heads_per_key_value_head : index - %is_first_subgroup = index.cmp eq, %subgroup, %c0 : index - %workitem_is_zero = index.cmp eq, %workitem, %c0 : index - %output_channel = index.add %workitem, %c0 : index - %partial_max_noalias, %partial_sum_noalias, %partial_output_noalias, %output_noalias = buffer.assume.noalias %partial_max, %partial_sum, %partial_output, %output : buffer, buffer, buffer, buffer - %partial_max_aligned = buffer.assume.alignment %partial_max_noalias {minimum_alignment = 16} : buffer - %partial_sum_aligned = buffer.assume.alignment %partial_sum_noalias {minimum_alignment = 16} : buffer - %partial_output_aligned = buffer.assume.alignment %partial_output_noalias {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %partial_max_view = buffer.view %partial_max_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - %partial_sum_view = buffer.view %partial_sum_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - %partial_output_view = buffer.view %partial_output_aligned[%c0_offset] : buffer -> view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_head_count]x128xf32> - %reduction_stage = buffer.alloca align(8) %reduction_stage_bytes : buffer - %reduction_stage_view = buffer.view %reduction_stage[%c0_offset] : buffer -> view<2xf32> - %lane_maximum = scf.if %is_first_subgroup -> (f32) { - %maximum = scf.for %block = [%lane to %key_value_block_count step %c64](%running_maximum = %negative_large : f32) -> (f32) { - %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 - %next_maximum = scalar.maxnumf %running_maximum, %block_maximum : f32 - scf.yield %next_maximum : f32 - } - scf.yield %maximum : f32 - } else { - scf.yield %negative_large : f32 - } - %subgroup_maximum = kernel.subgroup.reduce %lane_maximum : f32 - scf.if %workitem_is_zero { - view.store %subgroup_maximum, %reduction_stage_view[%c0] : f32, view<2xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %maximum = view.load %reduction_stage_view[%c0] : view<2xf32> -> f32 - %lane_sum = scf.if %is_first_subgroup -> (f32) { - %sum = scf.for %block = [%lane to %key_value_block_count step %c64](%running_sum = %c0_f32 : f32) -> (f32) { - %block_maximum = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 - %delta = scalar.subf %block_maximum, %maximum : f32 - %scale = scalar.expf %delta : f32 - view.store %scale, %partial_max_view[%key_value_head, %block, %query_row] : f32, view<[%key_value_head_count]x[%key_value_block_count]x16xf32> - %block_sum = view.load %partial_sum_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 - %scaled_sum = scalar.mulf %block_sum, %scale : f32 - %next_sum = scalar.addf %running_sum, %scaled_sum : f32 - scf.yield %next_sum : f32 - } - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %subgroup_sum = kernel.subgroup.reduce %lane_sum : f32 - scf.if %workitem_is_zero { - view.store %subgroup_sum, %reduction_stage_view[%c1] : f32, view<2xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %sum = view.load %reduction_stage_view[%c1] : view<2xf32> -> f32 - %unnormalized_output = scf.for %block = [%c0 to %key_value_block_count step %c1](%running_output = %c0_f32 : f32) -> (f32) unroll(%c4) schedule(interleaved) { - %scale = view.load %partial_max_view[%key_value_head, %block, %query_row] : view<[%key_value_head_count]x[%key_value_block_count]x16xf32> -> f32 - %block_output_f16 = view.load %partial_output_view[%key_value_head, %block, %query_row, %output_channel] : view<[%key_value_head_count]x[%key_value_block_count]x16x128xf16> -> f16 - %block_output = scalar.extf %block_output_f16 : f16 to f32 - %scaled_output = scalar.mulf %block_output, %scale : f32 - %next_output = scalar.addf %running_output, %scaled_output : f32 - scf.yield %next_output : f32 - } - %normalized_output = scalar.divf %unnormalized_output, %sum : f32 - view.store %normalized_output, %output_view[%query_head, %output_channel] : f32, view<[%query_head_count]x128xf32> - kernel.return -} - -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_case { - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<256x4x128xf16> - %value = check.generate.fill value(2.0) : tensor<256x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<256xf16> - %partial_max = check.generate.fill value(-1.0) : tensor<4x4x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<4x4x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<4x4x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<4xi32> - %output = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(2.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -// Only KV row zero participates. The finite F32 iota step overflows to F16 -// negative infinity at every later row, leaving blocks one through three -// entirely masked. Each empty split must publish the online-softmax identity -// instead of evaluating -inf - -inf and contaminating the final reduction. -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_masked_blocks_case { - %key_value_token_count = check.literal value(256) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<256x4x128xf16> - %value = check.generate.fill value(2.0) : tensor<256x4x128xf16> - %mask = check.generate.iota offset(0.0) step(-1e+30) : tensor<256xf16> - %partial_max = check.generate.fill value(-1.0) : tensor<4x4x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<4x4x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<4x4x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<4xi32> - %output = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(2.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<256x4x128xf16>, tensor<256x4x128xf16>, tensor<256xf16>, tensor<4x4x16xf32>, tensor<4x4x16xf32>, tensor<4x4x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -// The production Decode-513 shape combines eight full blocks with a one-row -// tail and selects the cooperative fused reducer. Constant nonzero V keeps the -// expected result exact while all four KV heads, 32 GQA heads, completion -// counters, partial tensors, and tail guards participate. A second invocation -// reuses the partial and counter storage with a different V tensor, making the -// completion-counter reset observable. -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_513_case { - %key_value_token_count = check.literal value(513) : index - %query = check.generate.fill value(1.0) : tensor<32x128xf32> - %key = check.generate.fill value(1.0) : tensor<513x4x128xf16> - %value0 = check.generate.fill value(2.0) : tensor<513x4x128xf16> - %value1 = check.generate.fill value(3.0) : tensor<513x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<513xf16> - %partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<4xi32> - %output0 = check.generate.fill value(-1.0) : tensor<32x128xf32> - %output1 = check.generate.fill value(-1.0) : tensor<32x128xf32> - %expected0 = check.generate.fill value(2.0) : tensor<32x128xf32> - %expected1 = check.generate.fill value(3.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output0) : [index](index, tensor<32x128xf32>, tensor<513x4x128xf16>, tensor<513x4x128xf16>, tensor<513xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output1) : [index](index, tensor<32x128xf32>, tensor<513x4x128xf16>, tensor<513x4x128xf16>, tensor<513xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - check.expect.close actual(%output0) expected(%expected0) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.expect.close actual(%output1) expected(%expected1) atol(0.001) rtol(0.001) nan(same) : tensor<32x128xf32> - check.return -} - -// Sixty-five KV rows force a second split containing one valid row. The mask -// selects that final row, whose iota values begin at 1024, so an omitted or -// uninitialized tail cannot accidentally satisfy the check. -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_tail_case { - %key_value_token_count = check.literal value(65) : index - %query = check.generate.fill value(1.0) : tensor<1x128xf32> - %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> - %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<65xf16> - %partial_max = check.generate.fill value(-1.0) : tensor<1x2x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<1x2x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<1x2x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %output = check.generate.fill value(-1.0) : tensor<1x128xf32> - %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<65xf16>, tensor<1x2x16xf32>, tensor<1x2x16xf32>, tensor<1x2x16x128xf16>, tensor<1xi32>, tensor<1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> - check.return -} - -// Thirty-two split-K blocks exercise the separate parallel reducer and the -// execution barrier between its producer and consumer dispatches. -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_case { - %key_value_token_count = check.literal value(2048) : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<2048x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<2048x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<2048xf16> - %partial_max = check.generate.fill value(-1.0) : tensor<4x32x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<4x32x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<4x32x16x128xf16> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : [index](index, tensor<32x128xf32>, tensor<2048x4x128xf16>, tensor<2048x4x128xf16>, tensor<2048xf16>, tensor<4x32x16xf32>, tensor<4x32x16xf32>, tensor<4x32x16x128xf16>) - kernel.launch @qwen3_moe_flash_attention_decode_split_reduce_f32[%key_value_token_count](%key_value_token_count, %partial_max, %partial_sum, %partial_output, %output) : [index](index, tensor<4x32x16xf32>, tensor<4x32x16xf32>, tensor<4x32x16x128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case { - %key_value_token_count = check.param.choice values([64, 65, 128, 256, 512, 513, 768, 1024, 1280, 2048]) name("key_value_token_count") : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> - // Reserve the bounded scratch capacity once; each specialization addresses - // only ceildiv(key_value_token_count, 64) blocks. - %partial_max = check.generate.fill value(-1.0) : tensor<4x512x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<4x512x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<4x512x16x128xf16> - %completion_counter = check.generate.fill value(0) : tensor<4xi32> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output, %completion_counter, %output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -// Long-context execution is one reusable producer/reducer command buffer. The -// harness records an explicit dispatch execution barrier between these calls -// and profiles their complete device-side span as one semantic operation. -check.case public @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case { - %key_value_token_count = check.param.choice values([2048, 32768]) name("key_value_token_count") : index - %query = check.generate.fill value(0.0) : tensor<32x128xf32> - %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<[%key_value_token_count]xf16> - %partial_max = check.generate.fill value(-1.0) : tensor<4x512x16xf32> - %partial_sum = check.generate.fill value(-1.0) : tensor<4x512x16xf32> - %partial_output = check.generate.fill value(-1.0) : tensor<4x512x16x128xf16> - %output = check.generate.fill value(1.0) : tensor<32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<32x128xf32> - kernel.launch @qwen3_moe_flash_attention_decode_split_produce_partials_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value, %mask, %partial_max, %partial_sum, %partial_output) : [index](index, tensor<32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]xf16>, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>) - kernel.launch @qwen3_moe_flash_attention_decode_split_reduce_f32[%key_value_token_count](%key_value_token_count, %partial_max, %partial_sum, %partial_output, %output) : [index](index, tensor<4x512x16xf32>, tensor<4x512x16xf32>, tensor<4x512x16x128xf16>, tensor<32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<32x128xf32> - check.return -} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_64 {key_value_token_count = 64} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_65 {key_value_token_count = 65} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_128 {key_value_token_count = 128} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_256 {key_value_token_count = 256} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_masked_blocks_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_256_masked_blocks - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_512 {key_value_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_513 {key_value_token_count = 513} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_768 {key_value_token_count = 768} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_1024 {key_value_token_count = 1024} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_1280 {key_value_token_count = 1280} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_2048_fused {key_value_token_count = 2048} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_2048 {key_value_token_count = 2048} - -check.benchmark<@qwen3_moe_flash_attention_decode_split_f32_f16_wmma_long_benchmark_case> @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_decode_32768 {key_value_token_count = 32768} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_next_q8_test.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_next_q8_test.loom deleted file mode 100644 index 04e658cd4fe1..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_decode_split_next_q8_test.loom +++ /dev/null @@ -1,113 +0,0 @@ -// Test-only differentials for fused decode-attention Q8 publication. The -// production kernels remain the actual providers; this module contributes only -// reference packing, masked metadata setup, and comparison cases. -target.decl @qwen3_moe_decode_split_gfx11_wave64 - -kernel.decl target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma(%key_value_token_count: index) launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer) - -kernel.decl target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8(%key_value_token_count: index) launch(%key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %partial_max: buffer, %partial_sum: buffer, %partial_output: buffer, %completion_counter: buffer, %output: buffer, %next_q8_output: buffer) - -// Publishes the canonical reference Q8_1 x4 layout for one 4096-element row. -kernel.def @qwen3_moe_flash_attention_decode_split_quantize_reference_4096() { - %c1 = index.constant 1 : index - %c32 = index.constant 32 : index - kernel.launch.config workgroups(%c32, %c1, %c1) workgroup_size(%c32, %c1, %c1) : index -} launch(%input: buffer, %output: buffer) { - %publish_output = scalar.constant true : i1 - %group_count = index.constant 32 : index - %group = kernel.workgroup.id : index - func.apply(%publish_output, %group_count, %group, %input, %output) : (i1, index, index, buffer, buffer) - kernel.return -} - -// Isolates the new publication phase for access-sanitized coverage without -// inflating the complete fused attention kernel beyond a short branch's range. -kernel.def target(@qwen3_moe_decode_split_gfx11_wave64) @qwen3_moe_flash_attention_decode_split_pack_completed_q8_test() { - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%c1, %c4, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%input: buffer, %output: buffer) { - %key_value_head = kernel.workgroup.id : index - func.apply(%key_value_head, %input, %output) : (index, buffer, buffer) - kernel.return -} - -// Models the request-owned Decode-513 mask inside its 576-row capacity class. -kernel.def @qwen3_moe_flash_attention_decode_split_mask_513_of_576() { - %c1 = index.constant 1 : index - %c64 = index.constant 64 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%mask: buffer) { - %workitem = kernel.workitem.id : index - %c63 = index.constant 63 : index - %c513 = index.constant 513 : index - %c0 = index.constant 0 : index - %c0_offset = index.constant 0 : offset - %negative_large = scalar.constant -1e+30 : f32 - %negative_infinity = scalar.fptrunc %negative_large : f32 to f16 - %valid = index.cmp ult, %workitem, %c63 : index - %row0 = index.add %c513, %workitem : index - %row = scf.select %valid, %row0, %c0 : index - %mask_view = buffer.view %mask[%c0_offset] : buffer -> view<576xf16> - scf.if %valid { - view.store %negative_infinity, %mask_view[%row] : f16, view<576xf16> - } - kernel.return -} - -// The fused producer must match the ordinary attention export followed by the -// canonical GGML packer for every F32 value and every packed byte. Two value -// tensors reuse both paths' partials and completion counters, making counter -// reset and stale-publication failures observable. -check.case public @qwen3_moe_flash_attention_decode_split_next_q8_capacity_576_differential_case { - %key_value_token_count = check.literal value(576) : index - %query_seed = check.param.seed base(0x514d4f4551465138) count(1) : i64 - %key_seed = check.param.seed base(0x514d4f454b465138) count(1) : i64 - %value_seed0 = check.param.seed base(0x514d4f4556465130) count(1) : i64 - %value_seed1 = check.param.seed base(0x514d4f4556465131) count(1) : i64 - %query = check.generate.random.uniform seed(%query_seed) range(-1.0 to 1.0) : tensor<32x128xf32> - %key = check.generate.random.uniform seed(%key_seed) range(-1.0 to 1.0) : tensor<576x4x128xf16> - %value0 = check.generate.random.uniform seed(%value_seed0) range(-1.0 to 1.0) : tensor<576x4x128xf16> - %value1 = check.generate.random.uniform seed(%value_seed1) range(-1.0 to 1.0) : tensor<576x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<576xf16> - %reference_partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> - %reference_partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> - %reference_partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> - %reference_counter = check.generate.fill value(0) : tensor<4xi32> - %actual_partial_max = check.generate.fill value(-1.0) : tensor<4x9x16xf32> - %actual_partial_sum = check.generate.fill value(-1.0) : tensor<4x9x16xf32> - %actual_partial_output = check.generate.fill value(-1.0) : tensor<4x9x16x128xf16> - %actual_counter = check.generate.fill value(0) : tensor<4xi32> - %reference_output0 = check.generate.fill value(-1.0) : tensor<32x128xf32> - %reference_output1 = check.generate.fill value(-1.0) : tensor<32x128xf32> - %actual_output0 = check.generate.fill value(-2.0) : tensor<32x128xf32> - %actual_output1 = check.generate.fill value(-2.0) : tensor<32x128xf32> - %reference_q8_0 = check.generate.fill value(0) : tensor<4608xi8> - %reference_q8_1 = check.generate.fill value(0) : tensor<4608xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<4608xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<4608xi8> - kernel.launch @qwen3_moe_flash_attention_decode_split_mask_513_of_576(%mask) : (tensor<576xf16>) - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %reference_partial_max, %reference_partial_sum, %reference_partial_output, %reference_counter, %reference_output0) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_decode_split_quantize_reference_4096(%reference_output0, %reference_q8_0) : (tensor<32x128xf32>, tensor<4608xi8>) - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value0, %mask, %actual_partial_max, %actual_partial_sum, %actual_partial_output, %actual_counter, %actual_output0, %actual_q8_0) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %reference_partial_max, %reference_partial_sum, %reference_partial_output, %reference_counter, %reference_output1) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_decode_split_quantize_reference_4096(%reference_output1, %reference_q8_1) : (tensor<32x128xf32>, tensor<4608xi8>) - kernel.launch @qwen3_moe_flash_attention_decode_split_f32_f16_wmma_next_q8[%key_value_token_count](%key_value_token_count, %query, %key, %value1, %mask, %actual_partial_max, %actual_partial_sum, %actual_partial_output, %actual_counter, %actual_output1, %actual_q8_1) : [index](index, tensor<32x128xf32>, tensor<576x4x128xf16>, tensor<576x4x128xf16>, tensor<576xf16>, tensor<4x9x16xf32>, tensor<4x9x16xf32>, tensor<4x9x16x128xf16>, tensor<4xi32>, tensor<32x128xf32>, tensor<4608xi8>) - check.expect.equal actual(%actual_output0) expected(%reference_output0) : tensor<32x128xf32> - check.expect.equal actual(%actual_q8_0) expected(%reference_q8_0) : tensor<4608xi8> - check.expect.equal actual(%actual_output1) expected(%reference_output1) : tensor<32x128xf32> - check.expect.equal actual(%actual_q8_1) expected(%reference_q8_1) : tensor<4608xi8> - check.return -} - -check.case public @qwen3_moe_flash_attention_decode_split_pack_completed_q8_differential_case { - %input_seed = check.param.seed base(0x514d4f4550385138) count(1) : i64 - %input = check.generate.random.uniform seed(%input_seed) range(-1.0 to 1.0) : tensor<32x128xf32> - %reference = check.generate.fill value(0) : tensor<4608xi8> - %actual = check.generate.fill value(1) : tensor<4608xi8> - kernel.launch @qwen3_moe_flash_attention_decode_split_quantize_reference_4096(%input, %reference) : (tensor<32x128xf32>, tensor<4608xi8>) - kernel.launch @qwen3_moe_flash_attention_decode_split_pack_completed_q8_test(%input, %actual) : (tensor<32x128xf32>, tensor<4608xi8>) - check.expect.equal actual(%actual) expected(%reference) : tensor<4608xi8> - check.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_f32_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_f32_f16_wmma.loom deleted file mode 100644 index d0b97ce7ee0f..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/flash_attention_f32_f16_wmma.loom +++ /dev/null @@ -1,870 +0,0 @@ -// Qwen3 MoE grouped-query FlashAttention for the F32-query/F16-cache path. -// -// One four-wave workgroup computes 16 query rows for one query head against -// 64 KV rows at a time. The ownership changes mirror the cooperative-matrix -// schedule used by llama.cpp's Vulkan CM1 kernel: -// -// 1. All workitems stage a scaled 16x128 F16 query tile. -// 2. Each wave computes one 16x16 QK score slice. -// 3. Scores cross LDS so each wave can normalize four complete query rows. -// 4. F16 probabilities cross LDS for four P*V WMMA steps. -// 5. Each active lane retains one four-channel F16 packet for each of its -// four query rows across subsequent 64-row KV blocks. -// -// K and V remain in the row-major llama.cpp cache layout -// [KV token][KV head][128]. Their aligned F16 fragments load directly from -// global memory; there is no expanded or repacked persistent allocation. QK -// and the online-softmax statistics remain F32, while the P*V accumulation and -// carried output match the Vulkan oracle's F16 policy. -amdgpu.target @qwen3_moe_attention_gfx11_wave64 {subgroup_size = 64} - -config.decl @qwen3_moe.attention.query_head_count : %value: index where [range(%value, 1, 64)] - -config.decl @qwen3_moe.attention.key_value_head_count : %value: index where [range(%value, 1, 64)] - -// Bounds the context copied by the test-only row-extraction kernel below. -config.decl @qwen3_moe.attention.test.context_capacity : %value: index where [range(%value, 1, 32768)] - -kernel.def target(@qwen3_moe_attention_gfx11_wave64) @qwen3_moe_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) { - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %c1 = index.constant 1 : index - %c15 = index.constant 15 : index - %c16 = index.constant 16 : index - %c256 = index.constant 256 : index - %padded_query_token_count = index.add %query_token_count, %c15 : index - %query_tile_count = index.div %padded_query_token_count, %c16 : index - kernel.launch.config workgroups(%query_tile_count, %query_head_count, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %output: buffer) where [range(%query_token_count, 1, 2048)] { - %bounded_key_value_token_count = index.assume %key_value_token_count [range(%key_value_token_count, 1, 32768)] : index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %key_value_head_count = config.get @qwen3_moe.attention.key_value_head_count : index - %query_tile = kernel.workgroup.id : index - %query_head0 = kernel.workgroup.id : index - %query_head, %launch_query_head_count = index.assume %query_head0, %query_head_count [lt(%query_head0, %query_head_count)] : index, index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %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 - %c2 = index.constant 2 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c15 = index.constant 15 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %query_stage_bytes = index.constant 4352 : offset - %score_stage_bytes = index.constant 6144 : offset - %probability_stage_bytes = index.constant 3072 : offset - %product_stage_bytes = index.constant 2048 : offset - %tail_key_value_stage_capacity = index.constant 8192 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %negative_large = scalar.constant -1e+30 : f32 - %head_size_f32 = scalar.constant 128.0 : f32 - %attention_scale = scalar.rsqrtf %head_size_f32 : f32 - %c0_f16 = scalar.constant 0.0 : f16 - %c1_f32 = scalar.constant 1.0 : f32 - %output_zero0 = vector.constant 0.0 : vector<4xf16> - %output_zero1 = vector.constant 0.0 : vector<4xf16> - %output_zero2 = vector.constant 0.0 : vector<4xf16> - %output_zero3 = vector.constant 0.0 : vector<4xf16> - %c0_f16x8 = vector.constant 0.0 : vector<8xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %negative_f32x4 = vector.constant -1e+30 : vector<4xf32> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %query_heads_per_key_value_head = index.div %query_head_count, %key_value_head_count : index - %key_value_head = index.div %query_head, %query_heads_per_key_value_head : index - %key_value_width = index.mul %key_value_head_count, %c128 : index - %key_value_head_base = index.mul %key_value_head, %c128 : index - %query_origin = index.mul %query_tile, %c16 : index - %full_key_value_block_count = index.div %bounded_key_value_token_count, %c64 : index - %full_key_value_token_count0 = index.mul %full_key_value_block_count, %c64 : index - %full_key_value_token_count = index.assume %full_key_value_token_count0 [range(%full_key_value_token_count0, 0, 32768), mul(%full_key_value_token_count0, 64)] : index - %tail_key_value_token_count = index.sub %bounded_key_value_token_count, %full_key_value_token_count : index - %has_key_value_tail = index.cmp ne, %tail_key_value_token_count, %c0 : index - %tail_key_value_stage_bytes = scf.select %has_key_value_tail, %tail_key_value_stage_capacity, %c0_offset : offset - %subgroup_score_column = index.mul %subgroup, %c16 : index - %subgroup_query_row = index.mul %subgroup, %c4 : index - %query_row0 = index.add %subgroup_query_row, %c0 : index - %query_row1 = index.add %subgroup_query_row, %c1 : index - %query_row2 = index.add %subgroup_query_row, %c2 : index - %query_row3 = index.add %subgroup_query_row, %c3 : index - %query_token0 = index.add %query_origin, %query_row0 : index - %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_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 - %lane_product_channel0 = index.rem %lane, %c16 : index - %lane_product_channel = index.mul %lane_product_channel0, %c4 : index - %lane_output_channel = index.mul %lane, %c4 : index - %lane_has_output = index.cmp ult, %lane, %c32 : index - // The padded LDS rows mirror the Vulkan oracle's ownership changes. Q uses - // eight spare F16 columns after its 128 channels. Score and probability - // transpose to key-major rows with eight spare columns after 16 queries. - // These strides avoid the bank pattern produced by dense transposed rows. - %query_transposed_layout = encoding.layout.strided [1, 136] : encoding - %probability_transposed_layout = encoding.layout.strided [1, 24] : encoding - %query_noalias, %key_noalias, %value_noalias, %mask_noalias, %output_noalias = buffer.assume.noalias %query, %key, %value, %mask, %output : buffer, buffer, buffer, buffer, buffer - %query_aligned = buffer.assume.alignment %query_noalias {minimum_alignment = 16} : buffer - %key_aligned = buffer.assume.alignment %key_noalias {minimum_alignment = 16} : buffer - %value_aligned = buffer.assume.alignment %value_noalias {minimum_alignment = 16} : buffer - %mask_aligned = buffer.assume.alignment %mask_noalias {minimum_alignment = 16} : buffer - %output_aligned = buffer.assume.alignment %output_noalias {minimum_alignment = 16} : buffer - %query_view = buffer.view %query_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x128xf32> - %key_view = buffer.view %key_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %value_view = buffer.view %value_aligned[%c0_offset] : buffer -> view<[%bounded_key_value_token_count]x[%key_value_width]xf16> - %mask_view = buffer.view %mask_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%bounded_key_value_token_count]xf16> - %output_view = buffer.view %output_aligned[%c0_offset] : buffer -> view<[%query_token_count]x[%query_head_count]x128xf32> - %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 - %product_stage = buffer.alloca align(16) %product_stage_bytes : buffer - %tail_key_value_stage = buffer.alloca align(16) %tail_key_value_stage_bytes : buffer - %query_stage_view = buffer.view %query_stage[%c0_offset] : buffer -> view<16x136xf16> - %query_transposed_view = buffer.view %query_stage[%c0_offset] : buffer -> view<128x16xf16, %query_transposed_layout> - %score_stage_view = buffer.view %score_stage[%c0_offset] : buffer -> view<64x24xf32> - %mask_summary_view = buffer.view %score_stage[%c0_offset] : buffer -> view<4xf32> - %probability_stage_view = buffer.view %probability_stage[%c0_offset] : buffer -> view<16x64xf16, %probability_transposed_layout> - %product_stage_view = buffer.view %product_stage[%c0_offset] : buffer -> view<16x64xf16> - %tail_key_value_stage_view = buffer.view %tail_key_value_stage[%c0_offset] : buffer -> view<32x128xf16> - // Scale and truncate Q exactly once. The Vulkan reference does this before - // entering its KV loop, making QK a native F16 WMMA while retaining F32 - // accumulation. - scf.for %load_iteration = [%c0 to %c8 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %local_query_row = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %query_token = index.add %query_origin, %local_query_row : index - %query_valid_load = index.cmp ult, %query_token, %query_token_count : index - %query_value = scf.if %query_valid_load -> (f16) { - %loaded = view.load %query_view[%query_token, %query_head, %query_channel] : view<[%query_token_count]x[%query_head_count]x128xf32> -> f32 - %scaled = scalar.mulf %loaded, %attention_scale : f32 - %truncated = scalar.fptrunc %scaled : f32 to f16 - scf.yield %truncated : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %query_value, %query_stage_view[%local_query_row, %query_channel] : f16, view<16x136xf16> - } - 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 - // KV blocks entirely at negative infinity; reducing the cached tile lets - // the workgroup skip QK, softmax, and P*V for those blocks. This mirrors - // the Vulkan oracle without requiring its auxiliary compact mask buffer. - %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 - 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 - 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 - 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 - scf.yield %value : f16 - } else { - scf.yield %c0_f16 : f16 - } - %mask_f16 = vector.from_elements %mask0, %mask1, %mask2, %mask3 : vector<4xf16> - %mask_summary_f32 = vector.extf %mask_f16 : vector<4xf16> to vector<4xf32> - %effective_mask = vector.select %query_valid, %mask_summary_f32, %negative_f32x4 : vector<4xf32> - %lane_mask_maximum = vector.reduce %effective_mask, %negative_large : vector<4xf32>, f32 - %subgroup_mask_maximum = kernel.subgroup.reduce %lane_mask_maximum : f32 - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_mask_maximum, %mask_summary_view[%subgroup] : f32, view<4xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %subgroup_mask_maxima = vector.load %mask_summary_view[%c0] : view<4xf32> -> vector<4xf32> - %workgroup_mask_maximum = vector.reduce %subgroup_mask_maxima, %negative_large : vector<4xf32>, f32 - %block_has_attention = scalar.cmpf ogt, %workgroup_mask_maximum, %negative_large : f32 - %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 = scf.if %block_has_attention -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - // Four independent wave-level WMMAs produce a 16x64 score tile. - %score_key_origin0 = index.add %key_origin, %subgroup_score_column : index - %last_full_key_tile_start = index.sub %bounded_key_value_token_count, %c15 : index - %score_key_origin = index.assume %score_key_origin0 [lt(%score_key_origin0, %last_full_key_tile_start)] : index - %score_init_values = vector.constant 0.0 : vector<4xf32> - %score_init = vector.fragment %score_init_values shape [%m, %n] : vector<4xf32> - %score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %score_init : vector<4xf32>) -> (vector<4xf32>) unroll schedule(recurrence) { - %key_channel = index.add %key_value_head_base, %head_tile : index - %key_fragment = vector.fragment.load %key_view[%score_key_origin, %key_channel] shape [%m, %k] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %score_fragment, %score_stage_view[%subgroup_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - // LDS transposes ownership from one 16-column score slice per wave to - // four complete query rows per wave. Every lane contributes one key - // column to each of those rows. - %raw_score0 = view.load %score_stage_view[%lane, %query_row0] : view<64x24xf32> -> f32 - %raw_score1 = view.load %score_stage_view[%lane, %query_row1] : view<64x24xf32> -> f32 - %raw_score2 = view.load %score_stage_view[%lane, %query_row2] : view<64x24xf32> -> f32 - %raw_score3 = view.load %score_stage_view[%lane, %query_row3] : view<64x24xf32> -> f32 - %mask_f32 = vector.extf %mask_f16 : vector<4xf16> to vector<4xf32> - %raw_scores = vector.from_elements %raw_score0, %raw_score1, %raw_score2, %raw_score3 : vector<4xf32> - %masked_scores0 = vector.addf %raw_scores, %mask_f32 : vector<4xf32> - %masked_scores = vector.select %query_valid, %masked_scores0, %negative_f32x4 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %query_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %old_delta = vector.subf %current_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Match the Vulkan CM1 ownership schedule by computing two sequential - // 64-channel output tiles. This keeps one P*V accumulator live per wave - // and reuses a 2 KiB exchange tile instead of retaining both halves. - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { - %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel0 = index.add %key_value_head_base, %output_tile_channel : index - %value_channel = index.add %value_channel0, %subgroup_product_channel : index - %product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %product_fragment = scf.for %key_tile = [%c0 to %c64 step %c16](%product_accumulator = %product_init : vector<8xf16>) -> (vector<8xf16>) unroll schedule(recurrence) { - %value_token0 = index.add %key_origin, %key_tile : index - %value_token = index.assume %value_token0 [lt(%value_token0, %last_full_key_tile_start)] : index - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %value_view[%value_token, %value_channel] shape [%k, %n] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> - } - vector.fragment.store %product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index - %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> - %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> - %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> - %updated_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> - scf.yield %updated_tile_output0, %updated_tile_output1, %updated_tile_output2, %updated_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // Complete every read before the next output tile overwrites LDS. - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %next_max, %next_sum, %next_output0, %next_output1, %next_output2, %next_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %current_max, %current_sum, %current_output0, %current_output1, %current_output2, %current_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - scf.yield %next_block_max, %next_block_sum, %next_block_output0, %next_block_output1, %next_block_output2, %next_block_output3 : vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - // A masked 32-row tile handles the final 1-63 KV rows with the same WMMA - // fragments as the aligned path. K and V are staged separately into product - // scratch, so every padded element is initialized and no physical padding is - // required of the caller. Each tile rounds probabilities to F16 before P*V. - %tail_score_wave = index.cmp ult, %subgroup, %c2 : index - %final_max, %final_sum, %final_output0, %final_output1, %final_output2, %final_output3 = scf.for %tail_key_origin = [%full_key_value_token_count to %bounded_key_value_token_count step %c32](%current_max = %full_max : vector<4xf32>, %current_sum = %full_sum : vector<4xf32>, %current_output0 = %full_output0 : vector<4xf16>, %current_output1 = %full_output1 : vector<4xf16>, %current_output2 = %full_output2 : vector<4xf16>, %current_output3 = %full_output3 : vector<4xf16>) -> (vector<4xf32>, vector<4xf32>, vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %tail_remaining = index.sub %bounded_key_value_token_count, %tail_key_origin : index - %tail_key_count = index.min %tail_remaining, %c32 : index - // Cooperatively stage one K tile, explicitly zeroing the padded rows. - scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %tail_key_row = index.div %linear, %c128 : index - %tail_key_channel = index.rem %linear, %c128 : index - %tail_key_valid = index.cmp ult, %tail_key_row, %tail_key_count : index - %tail_key_value = scf.if %tail_key_valid -> (f16) { - %tail_key_token0 = index.add %tail_key_origin, %tail_key_row : index - %tail_key_token = index.assume %tail_key_token0 [lt(%tail_key_token0, %bounded_key_value_token_count)] : index - %global_key_channel = index.add %key_value_head_base, %tail_key_channel : index - %loaded = view.load %key_view[%tail_key_token, %global_key_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %tail_key_value, %tail_key_value_stage_view[%tail_key_row, %tail_key_channel] : f16, view<32x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Waves zero and one compute the two 16-column QK fragments in this tile. - scf.if %tail_score_wave { - %tail_score_subgroup = index.assume %subgroup [range(%subgroup, 0, 1)] : index - %tail_score_column = index.mul %tail_score_subgroup, %c16 : index - %tail_score_init_values = vector.constant 0.0 : vector<4xf32> - %tail_score_init = vector.fragment %tail_score_init_values shape [%m, %n] : vector<4xf32> - %tail_score_fragment = scf.for %head_tile = [%c0 to %c128 step %c16](%score_accumulator = %tail_score_init : vector<4xf32>) -> (vector<4xf32>) unroll { - %key_fragment = vector.fragment.load %tail_key_value_stage_view[%tail_score_column, %head_tile] shape [%m, %k] : view<32x128xf16> -> vector<16xf16> - %query_fragment = vector.fragment.load %query_transposed_view[%head_tile, %c0] shape [%k, %n] : view<128x16xf16, %query_transposed_layout> -> vector<16xf16> - %next_score_accumulator = vector.mma %key_fragment, %query_fragment, %score_accumulator : vector<16xf16>, vector<16xf16>, vector<4xf32> - scf.yield %next_score_accumulator : vector<4xf32> - } - vector.fragment.store %tail_score_fragment, %score_stage_view[%tail_score_column, %c0] shape [%m, %n] : vector<4xf32>, view<64x24xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // LDS changes ownership from the QK wave to four query-row waves. Lanes - // beyond the logical tail never read the score or mask buffers. - %tail_lane_valid = index.cmp ult, %lane, %tail_key_count : index - %tail_valid0 = scalar.andi %tail_lane_valid, %query_valid0 : i1 - %tail_valid1 = scalar.andi %tail_lane_valid, %query_valid1 : i1 - %tail_valid2 = scalar.andi %tail_lane_valid, %query_valid2 : i1 - %tail_valid3 = scalar.andi %tail_lane_valid, %query_valid3 : i1 - %tail_valid = vector.from_elements %tail_valid0, %tail_valid1, %tail_valid2, %tail_valid3 : vector<4xi1> - %tail_key_token0 = index.add %tail_key_origin, %lane : index - %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<64x24xf32> -> 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_f32 = scalar.extf %mask_f16 : f16 to f32 - %score = scalar.addf %raw_score, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %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<64x24xf32> -> 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_f32 = scalar.extf %mask_f16 : f16 to f32 - %score = scalar.addf %raw_score, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %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<64x24xf32> -> 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_f32 = scalar.extf %mask_f16 : f16 to f32 - %score = scalar.addf %raw_score, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %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<64x24xf32> -> 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_f32 = scalar.extf %mask_f16 : f16 to f32 - %score = scalar.addf %raw_score, %mask_f32 : f32 - scf.yield %score : f32 - } else { - scf.yield %negative_large : f32 - } - %masked_scores = vector.from_elements %masked_score0, %masked_score1, %masked_score2, %masked_score3 : vector<4xf32> - %block_max = kernel.subgroup.reduce %masked_scores : vector<4xf32> - %next_max = vector.maxnumf %current_max, %block_max : vector<4xf32> - %score_delta = vector.subf %masked_scores, %next_max : vector<4xf32> - %raw_probability = vector.expf %score_delta : vector<4xf32> - %probability = vector.select %tail_valid, %raw_probability, %c0_f32x4 : vector<4xf32> - %block_sum = kernel.subgroup.reduce %probability : vector<4xf32> - %old_delta = vector.subf %current_max, %next_max : vector<4xf32> - %old_scale = vector.expf %old_delta : vector<4xf32> - %scaled_current_sum = vector.mulf %current_sum, %old_scale : vector<4xf32> - %next_sum = vector.addf %scaled_current_sum, %block_sum : vector<4xf32> - %probability_f16 = vector.fptrunc %probability : vector<4xf32> to vector<4xf16> - %probability0 = vector.extract %probability_f16[0] : vector<4xf16> -> f16 - %probability1 = vector.extract %probability_f16[1] : vector<4xf16> -> f16 - %probability2 = vector.extract %probability_f16[2] : vector<4xf16> -> f16 - %probability3 = vector.extract %probability_f16[3] : vector<4xf16> -> f16 - view.store %probability0, %probability_stage_view[%query_row0, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability1, %probability_stage_view[%query_row1, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability2, %probability_stage_view[%query_row2, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - view.store %probability3, %probability_stage_view[%query_row3, %lane] : f16, view<16x64xf16, %probability_transposed_layout> - kernel.barrier scope(workgroup) ordering(acq_rel) - // Reuse product scratch for V after every probability is resident in its - // disjoint LDS tile. - scf.for %load_iteration = [%c0 to %c16 step %c1] unroll { - %linear = index.madd %load_iteration, %c256, %workitem : index - %tail_value_row = index.div %linear, %c128 : index - %tail_value_channel = index.rem %linear, %c128 : index - %tail_value_valid = index.cmp ult, %tail_value_row, %tail_key_count : index - %tail_value = scf.if %tail_value_valid -> (f16) { - %tail_value_token0 = index.add %tail_key_origin, %tail_value_row : index - %tail_value_token = index.assume %tail_value_token0 [lt(%tail_value_token0, %bounded_key_value_token_count)] : index - %global_value_channel = index.add %key_value_head_base, %tail_value_channel : index - %loaded = view.load %value_view[%tail_value_token, %global_value_channel] : view<[%bounded_key_value_token_count]x[%key_value_width]xf16> -> f16 - scf.yield %loaded : f16 - } else { - scf.yield %c0_f16 : f16 - } - view.store %tail_value, %tail_key_value_stage_view[%tail_value_row, %tail_value_channel] : f16, view<32x128xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - // Keep tail V staging separate from the 2 KiB product exchange, then use - // the same two 64-channel phases as the aligned path. - %old_scale_f16 = vector.fptrunc %old_scale : vector<4xf32> to vector<4xf16> - %old_scale0_scalar = vector.extract %old_scale_f16[0] : vector<4xf16> -> f16 - %old_scale1_scalar = vector.extract %old_scale_f16[1] : vector<4xf16> -> f16 - %old_scale2_scalar = vector.extract %old_scale_f16[2] : vector<4xf16> -> f16 - %old_scale3_scalar = vector.extract %old_scale_f16[3] : vector<4xf16> -> f16 - %old_scale0 = vector.splat %old_scale0_scalar : vector<4xf16> - %old_scale1 = vector.splat %old_scale1_scalar : vector<4xf16> - %old_scale2 = vector.splat %old_scale2_scalar : vector<4xf16> - %old_scale3 = vector.splat %old_scale3_scalar : vector<4xf16> - %scaled_current_output0 = vector.mulf %current_output0, %old_scale0 : vector<4xf16> - %scaled_current_output1 = vector.mulf %current_output1, %old_scale1 : vector<4xf16> - %scaled_current_output2 = vector.mulf %current_output2, %old_scale2 : vector<4xf16> - %scaled_current_output3 = vector.mulf %current_output3, %old_scale3 : vector<4xf16> - %next_output0, %next_output1, %next_output2, %next_output3 = scf.for %output_tile = [%c0 to %c2 step %c1](%tile_output0 = %scaled_current_output0 : vector<4xf16>, %tile_output1 = %scaled_current_output1 : vector<4xf16>, %tile_output2 = %scaled_current_output2 : vector<4xf16>, %tile_output3 = %scaled_current_output3 : vector<4xf16>) -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) unroll { - %output_tile_channel = index.mul %output_tile, %c64 : index - %value_channel = index.add %output_tile_channel, %subgroup_product_channel : index - %tail_product_init = vector.fragment %c0_f16x8 shape [%m, %n] : vector<8xf16> - %tail_product_fragment = scf.for %key_tile = [%c0 to %c32 step %c16](%product_accumulator = %tail_product_init : vector<8xf16>) -> (vector<8xf16>) unroll { - %probability_fragment = vector.fragment.load %probability_stage_view[%c0, %key_tile] shape [%m, %k] : view<16x64xf16, %probability_transposed_layout> -> vector<16xf16> - %value_fragment = vector.fragment.load %tail_key_value_stage_view[%key_tile, %value_channel] shape [%k, %n] : view<32x128xf16> -> vector<16xf16> - %next_product_accumulator = vector.mma %probability_fragment, %value_fragment, %product_accumulator : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %next_product_accumulator : vector<8xf16> - } - vector.fragment.store %tail_product_fragment, %product_stage_view[%c0, %subgroup_product_channel] shape [%m, %n] : vector<8xf16>, view<16x64xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %owns_output_tile = index.cmp eq, %lane_output_tile, %output_tile : index - %updated_output0, %updated_output1, %updated_output2, %updated_output3 = scf.if %owns_output_tile -> (vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16>) { - %block_output0 = vector.load %product_stage_view[%query_row0, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output1 = vector.load %product_stage_view[%query_row1, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output2 = vector.load %product_stage_view[%query_row2, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %block_output3 = vector.load %product_stage_view[%query_row3, %lane_product_channel] : view<16x64xf16> -> vector<4xf16> - %updated_tile_output0 = vector.addf %tile_output0, %block_output0 : vector<4xf16> - %updated_tile_output1 = vector.addf %tile_output1, %block_output1 : vector<4xf16> - %updated_tile_output2 = vector.addf %tile_output2, %block_output2 : vector<4xf16> - %updated_tile_output3 = vector.addf %tile_output3, %block_output3 : vector<4xf16> - scf.yield %updated_tile_output0, %updated_tile_output1, %updated_tile_output2, %updated_tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } else { - scf.yield %tile_output0, %tile_output1, %tile_output2, %tile_output3 : vector<4xf16>, vector<4xf16>, vector<4xf16>, vector<4xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %updated_output0, %updated_output1, %updated_output2, %updated_output3 : 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> - } - // Normalize and publish the lane-owned packets. Query-tail rows never - // participate in the output store. - scf.if %lane_has_output { - %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 - %sum3_scalar = vector.extract %final_sum[3] : vector<4xf32> -> f32 - %inverse_sum0_f32 = scalar.divf %c1_f32, %sum0_scalar : f32 - %inverse_sum1_f32 = scalar.divf %c1_f32, %sum1_scalar : f32 - %inverse_sum2_f32 = scalar.divf %c1_f32, %sum2_scalar : f32 - %inverse_sum3_f32 = scalar.divf %c1_f32, %sum3_scalar : f32 - %inverse_sum0_f16 = scalar.fptrunc %inverse_sum0_f32 : f32 to f16 - %inverse_sum1_f16 = scalar.fptrunc %inverse_sum1_f32 : f32 to f16 - %inverse_sum2_f16 = scalar.fptrunc %inverse_sum2_f32 : f32 to f16 - %inverse_sum3_f16 = scalar.fptrunc %inverse_sum3_f32 : f32 to f16 - %inverse_sum0 = vector.splat %inverse_sum0_f16 : vector<4xf16> - %inverse_sum1 = vector.splat %inverse_sum1_f16 : vector<4xf16> - %inverse_sum2 = vector.splat %inverse_sum2_f16 : vector<4xf16> - %inverse_sum3 = vector.splat %inverse_sum3_f16 : vector<4xf16> - %normalized0_f16 = vector.mulf %final_output0, %inverse_sum0 : vector<4xf16> - %normalized1_f16 = vector.mulf %final_output1, %inverse_sum1 : vector<4xf16> - %normalized2_f16 = vector.mulf %final_output2, %inverse_sum2 : vector<4xf16> - %normalized3_f16 = vector.mulf %final_output3, %inverse_sum3 : vector<4xf16> - %normalized0 = vector.extf %normalized0_f16 : vector<4xf16> to vector<4xf32> - %normalized1 = vector.extf %normalized1_f16 : vector<4xf16> to vector<4xf32> - %normalized2 = vector.extf %normalized2_f16 : vector<4xf16> to vector<4xf32> - %normalized3 = vector.extf %normalized3_f16 : vector<4xf16> to vector<4xf32> - scf.if %query_valid0 { - vector.store %normalized0, %output_view[%query_token0, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> - } - scf.if %query_valid1 { - vector.store %normalized1, %output_view[%query_token1, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> - } - scf.if %query_valid2 { - vector.store %normalized2, %output_view[%query_token2, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> - } - scf.if %query_valid3 { - vector.store %normalized3, %output_view[%query_token3, %query_head, %lane_output_channel] : vector<4xf32>, view<[%query_token_count]x[%query_head_count]x128xf32> - } - } - kernel.return -} - -// Test-only causal-mask construction for the production 14-token witness. -// Finite F16 minima retain exact zero probabilities without requiring infinity -// support from synthetic tensor generators. Selective linking drops this -// helper from production roots. -kernel.def target(@qwen3_moe_attention_gfx11_wave64) @qwen3_moe_flash_attention_test_make_causal_mask() { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%mask: buffer) { - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c14 = index.constant 14 : index - %c196 = index.constant 196 : index - %c0_offset = index.constant 0 : offset - %c0_f16 = scalar.constant 0.0 : f16 - %negative_f16 = scalar.constant -65504.0 : f16 - %workitem = kernel.workitem.id : index - %in_bounds = index.cmp ult, %workitem, %c196 : index - %mask_noalias = buffer.assume.noalias %mask : buffer - %mask_view = buffer.view %mask_noalias[%c0_offset] : buffer -> view<14x14xf16> - scf.if %in_bounds { - %row = index.div %workitem, %c14 : index - %column = index.rem %workitem, %c14 : index - %row_limit = index.add %row, %c1 : index - %is_visible = index.cmp ult, %column, %row_limit : index - %value = scf.select %is_visible, %c0_f16, %negative_f16 : f16 - view.store %value, %mask_view[%row, %column] : f16, view<14x14xf16> - } - kernel.return -} - -// Test-only row extraction for comparing one multirow attention dispatch with -// independent one-row dispatches over identical data. Selective linking drops -// this helper from production roots. -kernel.def target(@qwen3_moe_attention_gfx11_wave64) @qwen3_moe_flash_attention_test_extract_row(%query_token_count: index, %context_count: index, %source_row: index) { - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %context_capacity = config.get @qwen3_moe.attention.test.context_capacity : index - %c128 = index.constant 128 : index - %c255 = index.constant 255 : index - %c256 = index.constant 256 : index - %c1 = index.constant 1 : index - %query_element_count = index.mul %query_head_count, %c128 : index - %element_count = index.add %query_element_count, %context_capacity : index - %padded_element_count = index.add %element_count, %c255 : index - %workgroup_count = index.div %padded_element_count, %c256 : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%query_token_count: index, %context_count: index, %source_row: index, %source_query: buffer, %source_mask: buffer, %target_query: buffer, %target_mask: buffer) { - %c0 = index.constant 0 : index - %c0_offset = index.constant 0 : offset - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %query_head_count = config.get @qwen3_moe.attention.query_head_count : index - %context_capacity = config.get @qwen3_moe.attention.test.context_capacity : index - %bounded_query_token_count = index.assume %query_token_count [range(%query_token_count, 1, 2048)] : index - %bounded_context_count = index.assume %context_count [range(%context_count, 1, 32768), le(%context_count, %context_capacity)] : index - %bounded_source_row, %source_query_token_count = index.assume %source_row, %bounded_query_token_count [range(%source_row, 0, 2047), lt(%source_row, %bounded_query_token_count)] : index, index - %workgroup = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %linear = index.madd %workgroup, %c256, %workitem : index - %query_element_count = index.mul %query_head_count, %c128 : index - %element_count = index.add %query_element_count, %bounded_context_count : index - %in_bounds = index.cmp ult, %linear, %element_count : index - %source_query_noalias, %source_mask_noalias, %target_query_noalias, %target_mask_noalias = buffer.assume.noalias %source_query, %source_mask, %target_query, %target_mask : buffer, buffer, buffer, buffer - %source_query_view = buffer.view %source_query_noalias[%c0_offset] : buffer -> view<[%source_query_token_count]x[%query_head_count]x128xf32> - %source_mask_view = buffer.view %source_mask_noalias[%c0_offset] : buffer -> view<[%source_query_token_count]x[%bounded_context_count]xf16> - %target_query_view = buffer.view %target_query_noalias[%c0_offset] : buffer -> view<1x[%query_head_count]x128xf32> - %target_mask_view = buffer.view %target_mask_noalias[%c0_offset] : buffer -> view<1x[%bounded_context_count]xf16> - scf.if %in_bounds { - %is_query_element = index.cmp ult, %linear, %query_element_count : index - scf.if %is_query_element { - %query_head = index.div %linear, %c128 : index - %query_channel = index.rem %linear, %c128 : index - %value = view.load %source_query_view[%bounded_source_row, %query_head, %query_channel] : view<[%source_query_token_count]x[%query_head_count]x128xf32> -> f32 - view.store %value, %target_query_view[%c0, %query_head, %query_channel] : f32, view<1x[%query_head_count]x128xf32> - } else { - %mask_column = index.sub %linear, %query_element_count : index - %value = view.load %source_mask_view[%bounded_source_row, %mask_column] : view<[%source_query_token_count]x[%bounded_context_count]xf16> -> f16 - view.store %value, %target_mask_view[%c0, %mask_column] : f16, view<1x[%bounded_context_count]xf16> - } - } - kernel.return -} - -// The mask selects the first KV row exactly. QK, F16 probability conversion, -// GQA addressing, and the P*V path all execute, while the expected result is -// the first V row and remains auditable as an iota. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_selected_row_case { - %query_token_count = check.literal value(1) : index - %key_value_token_count = check.literal value(64) : index - %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> - %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<64x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-10000.0) : tensor<1x64xf16> - %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x1x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<1x64xf16>, tensor<1x1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> - check.return -} - -// Four query rows select four distinct KV rows. The 33-element mask period -// maps row-major index row*32+column to zero exactly when row == column for -// rows zero through three. The expected output is therefore the first four -// rows of V, preserving an auditable iota while distinguishing every component -// in the four-row per-wave ownership path. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_multirow_selected_case { - %query_token_count = check.literal value(4) : index - %key_value_token_count = check.literal value(32) : index - %query = check.generate.fill value(1.0) : tensor<4x1x128xf32> - %key = check.generate.fill value(1.0) : tensor<32x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<32x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-10000.0) period(33) : tensor<4x32xf16> - %output = check.generate.fill value(-1.0) : tensor<4x1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<4x1x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x1x128xf32>, tensor<32x1x128xf16>, tensor<32x1x128xf16>, tensor<4x32xf16>, tensor<4x1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<4x1x128xf32> - check.return -} - -// A unit mask step gives every row the same geometric softmax distribution, -// shifted to begin at the row-matched KV entry. Values advance by 1/4096, so -// each KV row adds exactly 1/32 and the infinite-series weighted row offset is -// 1/(32*(e-1)). Terms that wrap at period 33 are below F16 significance. This -// exercises four independent online-softmax and P*V states while retaining a -// compact closed-form expected iota. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_multirow_online_case { - %query_token_count = check.literal value(4) : index - %key_value_token_count = check.literal value(32) : index - %query = check.generate.fill value(1.0) : tensor<4x1x128xf32> - %key = check.generate.fill value(1.0) : tensor<32x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.000244140625) : tensor<32x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-1.0) period(33) : tensor<4x32xf16> - %output = check.generate.fill value(-1.0) : tensor<4x1x128xf32> - %expected = check.generate.iota offset(0.01818677224124075) step(0.000244140625) : tensor<4x1x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x1x128xf32>, tensor<32x1x128xf16>, tensor<32x1x128xf16>, tensor<4x32xf16>, tensor<4x1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<4x1x128xf32> - check.return -} - -// Compare the first four rows of one fourteen-row launch with four one-row -// launches over identical, nonuniform data and the production causal mask. -// Fourteen rows retain the exact dynamic dense-row stride, ownership, and -// query-validity pressure, while the one-row path is the proven containment. -// Row extraction only reshapes bindings and performs no attention arithmetic. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_multirow_differential_case { - %query_token_count = check.literal value(14) : index - %c1 = check.literal value(1) : index - %key_value_token_count = check.literal value(14) : index - %row0 = check.literal value(0) : index - %row1 = check.literal value(1) : index - %row2 = check.literal value(2) : index - %row3 = check.literal value(3) : index - %query_seed = check.param.seed base(0x514d4f4551554552) count(1) : i64 - %key_seed = check.param.seed base(0x514d4f454b455931) count(1) : i64 - %value_seed = check.param.seed base(0x514d4f4556414c31) count(1) : i64 - %query = check.generate.random.uniform seed(%query_seed) range(-1.0 to 1.0) : tensor<14x32x128xf32> - %key = check.generate.random.uniform seed(%key_seed) range(-1.0 to 1.0) : tensor<14x4x128xf16> - %value = check.generate.random.uniform seed(%value_seed) range(-1.0 to 1.0) : tensor<14x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<14x14xf16> - %actual = check.generate.fill value(-1.0) : tensor<14x32x128xf32> - %row_query = check.generate.fill value(0.0) : tensor<1x32x128xf32> - %row_mask = check.generate.fill value(0.0) : tensor<1x14xf16> - %expected0 = check.generate.fill value(0.0) : tensor<1x32x128xf32> - %expected1 = check.generate.fill value(0.0) : tensor<1x32x128xf32> - %expected2 = check.generate.fill value(0.0) : tensor<1x32x128xf32> - %expected3 = check.generate.fill value(0.0) : tensor<1x32x128xf32> - %actual0 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> - %actual1 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> - %actual2 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> - %actual3 = check.generate.fill value(-1.0) : tensor<1x32x128xf32> - kernel.launch @qwen3_moe_flash_attention_test_make_causal_mask(%mask) : (tensor<14x14xf16>) - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %actual) : [index, index](index, index, tensor<14x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<14x14xf16>, tensor<14x32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row0](%query_token_count, %key_value_token_count, %row0, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected0) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row0](%query_token_count, %key_value_token_count, %row0, %actual, %mask, %actual0, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row1](%query_token_count, %key_value_token_count, %row1, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected1) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row1](%query_token_count, %key_value_token_count, %row1, %actual, %mask, %actual1, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row2](%query_token_count, %key_value_token_count, %row2, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected2) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row2](%query_token_count, %key_value_token_count, %row2, %actual, %mask, %actual2, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row3](%query_token_count, %key_value_token_count, %row3, %query, %mask, %row_query, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%c1, %key_value_token_count](%c1, %key_value_token_count, %row_query, %key, %value, %row_mask, %expected3) : [index, index](index, index, tensor<1x32x128xf32>, tensor<14x4x128xf16>, tensor<14x4x128xf16>, tensor<1x14xf16>, tensor<1x32x128xf32>) - kernel.launch @qwen3_moe_flash_attention_test_extract_row[%query_token_count, %key_value_token_count, %row3](%query_token_count, %key_value_token_count, %row3, %actual, %mask, %actual3, %row_mask) : [index, index, index](index, index, index, tensor<14x32x128xf32>, tensor<14x14xf16>, tensor<1x32x128xf32>, tensor<1x14xf16>) - check.expect.close actual(%actual0) expected(%expected0) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> - check.expect.close actual(%actual1) expected(%expected1) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> - check.expect.close actual(%actual2) expected(%expected2) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> - check.expect.close actual(%actual3) expected(%expected3) atol(0.001) rtol(0.001) nan(same) : tensor<1x32x128xf32> - check.return -} - -// Seventeen rows force a partial second query tile. Equal scores and constant -// values make every valid output exactly two while still exercising online -// normalization and the query-tail guards. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_query_tail_case { - %query_token_count = check.literal value(17) : index - %key_value_token_count = check.literal value(64) : index - %query = check.generate.fill value(1.0) : tensor<17x1x128xf32> - %key = check.generate.fill value(1.0) : tensor<64x1x128xf16> - %value = check.generate.fill value(2.0) : tensor<64x1x128xf16> - %mask = check.generate.fill value(0.0) : tensor<17x64xf16> - %output = check.generate.fill value(-1.0) : tensor<17x1x128xf32> - %expected = check.generate.fill value(2.0) : tensor<17x1x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<17x1x128xf32>, tensor<64x1x128xf16>, tensor<64x1x128xf16>, tensor<17x64xf16>, tensor<17x1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<17x1x128xf32> - check.return -} - -// Sixty-five KV rows force one masked cleanup tile after a full WMMA block. -// The mask selects only that final row, whose iota values begin at 1024, so -// omitting the cleanup path cannot accidentally satisfy the check. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_key_value_tail_case { - %query_token_count = check.literal value(1) : index - %key_value_token_count = check.literal value(65) : index - %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> - %key = check.generate.fill value(1.0) : tensor<65x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<65x1x128xf16> - %mask = check.generate.iota offset(-64000.0) step(1000.0) : tensor<1x65xf16> - %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> - %expected = check.generate.iota offset(1024.0) step(0.125) : tensor<1x1x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<65x1x128xf16>, tensor<65x1x128xf16>, tensor<1x65xf16>, tensor<1x1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> - check.return -} - -// The first 64-row block contains the selected value while every mask in the -// second block rounds to negative infinity. This preserves a finite online -// softmax state while exercising the workgroup-wide block-pruning path. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_pruned_block_case { - %query_token_count = check.literal value(1) : index - %key_value_token_count = check.literal value(128) : index - %query = check.generate.fill value(1.0) : tensor<1x1x128xf32> - %key = check.generate.fill value(1.0) : tensor<128x1x128xf16> - %value = check.generate.iota offset(0.0) step(0.125) : tensor<128x1x128xf16> - %mask = check.generate.iota offset(0.0) step(-2000.0) : tensor<1x128xf16> - %output = check.generate.fill value(-1.0) : tensor<1x1x128xf32> - %expected = check.generate.iota offset(0.0) step(0.125) : tensor<1x1x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x1x128xf32>, tensor<128x1x128xf16>, tensor<128x1x128xf16>, tensor<1x128xf16>, tensor<1x1x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.001) rtol(0.001) nan(same) : tensor<1x1x128xf32> - check.return -} - -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case { - %query_token_count = check.param.choice values([1, 32, 64, 128, 192, 255, 256, 257, 384, 511, 512, 513, 768, 1023, 1024, 1025, 1280, 1536, 1792, 2048]) name("query_token_count") : index - %key_value_token_count = check.param.choice values([64, 128, 192, 255, 256, 257, 384, 511, 512, 513, 768, 1023, 1024, 1025, 1280, 1536, 1792, 2048, 32768]) name("key_value_token_count") : index - %query = check.generate.fill value(0.0) : tensor<[%query_token_count]x32x128xf32> - %key = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<[%key_value_token_count]x4x128xf16> - %mask = check.generate.fill value(0.0) : tensor<[%query_token_count]x[%key_value_token_count]xf16> - %output = check.generate.fill value(1.0) : tensor<[%query_token_count]x32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<[%query_token_count]x32x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<[%query_token_count]x32x128xf32>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%key_value_token_count]x4x128xf16>, tensor<[%query_token_count]x[%key_value_token_count]xf16>, tensor<[%query_token_count]x32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%query_token_count]x32x128xf32> - check.return -} - -// Four finite 64-row blocks followed by four negative-infinity blocks model -// the aggregate computed/pruned work ratio of 512-token causal prefill while -// keeping every workgroup's path identical for a stable microbenchmark. -check.case public @qwen3_moe_flash_attention_f32_f16_wmma_pruned_half_benchmark_case { - %query_token_count = check.literal value(512) : index - %key_value_token_count = check.literal value(512) : index - %query = check.generate.fill value(0.0) : tensor<512x32x128xf32> - %key = check.generate.fill value(0.0) : tensor<512x4x128xf16> - %value = check.generate.fill value(0.0) : tensor<512x4x128xf16> - %mask = check.generate.iota offset(0.0) step(-300.0) period(512) : tensor<512x512xf16> - %output = check.generate.fill value(1.0) : tensor<512x32x128xf32> - %expected = check.generate.fill value(0.0) : tensor<512x32x128xf32> - kernel.launch @qwen3_moe_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<512x32x128xf32>, tensor<512x4x128xf16>, tensor<512x4x128xf16>, tensor<512x512xf16>, tensor<512x32x128xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<512x32x128xf32> - check.return -} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_selected_row_case> @qwen3_moe_flash_attention_f32_f16_wmma_selected_row - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_query_tail_case> @qwen3_moe_flash_attention_f32_f16_wmma_query_tail - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_key_value_tail_case> @qwen3_moe_flash_attention_f32_f16_wmma_key_value_tail - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_pruned_block_case> @qwen3_moe_flash_attention_f32_f16_wmma_pruned_block - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_pruned_half_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_pruned_half - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_decode_256 {key_value_token_count = 256, query_token_count = 1} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_decode_2048 {key_value_token_count = 2048, query_token_count = 1} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_decode_32768 {key_value_token_count = 32768, query_token_count = 1} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_32 {key_value_token_count = 256, query_token_count = 32} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_128 {key_value_token_count = 256, query_token_count = 128} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512 {key_value_token_count = 512, query_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_context_1024 {key_value_token_count = 1024, query_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_context_1536 {key_value_token_count = 1536, query_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_512_context_2048 {key_value_token_count = 2048, query_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_1024 {key_value_token_count = 1024, query_token_count = 1024} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_prefill_2048 {key_value_token_count = 2048, query_token_count = 2048} - -// These aligned and boundary-adjacent self-attention shapes expose launch or -// tail cliffs that a powers-of-two-only benchmark would hide. They are -// measurement witnesses for one shape-specialized kernel, not routing buckets. -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_64 {key_value_token_count = 64, query_token_count = 64} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_128 {key_value_token_count = 128, query_token_count = 128} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_192 {key_value_token_count = 192, query_token_count = 192} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_255 {key_value_token_count = 255, query_token_count = 255} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_256 {key_value_token_count = 256, query_token_count = 256} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_257 {key_value_token_count = 257, query_token_count = 257} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_384 {key_value_token_count = 384, query_token_count = 384} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_511 {key_value_token_count = 511, query_token_count = 511} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_512 {key_value_token_count = 512, query_token_count = 512} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_513 {key_value_token_count = 513, query_token_count = 513} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_768 {key_value_token_count = 768, query_token_count = 768} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1023 {key_value_token_count = 1023, query_token_count = 1023} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1024 {key_value_token_count = 1024, query_token_count = 1024} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1025 {key_value_token_count = 1025, query_token_count = 1025} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1280 {key_value_token_count = 1280, query_token_count = 1280} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1536 {key_value_token_count = 1536, query_token_count = 1536} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_1792 {key_value_token_count = 1792, query_token_count = 1792} - -check.benchmark<@qwen3_moe_flash_attention_f32_f16_wmma_benchmark_case> @qwen3_moe_flash_attention_f32_f16_wmma_self_2048 {key_value_token_count = 2048, query_token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_prepare_target.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_prepare_target.loom new file mode 100644 index 000000000000..feb0e4fb8a6f --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_prepare_target.loom @@ -0,0 +1,4 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @qwen3_moe_attention_prepare_gfx11_wave32 {subgroup_size = 32} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom new file mode 100644 index 000000000000..d7f4d82b1a8e --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/attention_qkv_postprocess_fused_decode.loom @@ -0,0 +1,73 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.attention.qkv_postprocess_fused_decode.body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) + +template.decl @qwen3_moe.attention.qkv_quantized.body(%value_uses_q6_index: index, %publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %query_output: buffer, %key_output: buffer, %value_output: buffer) + +template.decl @qwen3_moe.attention.postprocess.head_body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: buffer, %arg6: buffer, %arg7: buffer, %arg8: buffer, %arg9: buffer, %arg10: buffer, %arg11: buffer, %arg12: buffer, %arg13: buffer, %arg14: buffer, %arg15: buffer, %arg16: buffer) + +config.decl @qwen3_moe.attention.query_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.key_value_size : %value: index where [range(%value, 1, 262144)] + +config.decl @qwen3_moe.attention.head_size : %value: index where [range(%value, 128, 128)] + +// The fused export sources the value storage format from model config. +template.def<@qwen3_moe.attention.qkv_postprocess_fused_decode.body> device @qwen3_moe_attention_qkv_postprocess_fused_decode_body(%value_uses_q6_index: index, %token_count: index, %cache_row_count: index, %q8_input: buffer, %query_weight: buffer, %key_weight: buffer, %value_weight: buffer, %positions: buffer, %key_cache_indices: buffer, %value_cache_indices: buffer, %query_output_raw: buffer, %key_output_raw: buffer, %value_output_raw: buffer, %query_norm_weight: buffer, %key_norm_weight: buffer, %inverse_frequencies: buffer, %query_output: buffer, %key_cache: buffer, %value_cache: buffer, %completion_counters: buffer) { + %publish_projection_output = scalar.constant true : i1 + %projection_token = kernel.workgroup.id : index + template.apply<@qwen3_moe.attention.qkv_quantized.body>(%value_uses_q6_index, %publish_projection_output, %token_count, %projection_token, %q8_input, %query_weight, %key_weight, %value_weight, %query_output_raw, %key_output_raw, %value_output_raw) : (index, i1, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + %query_size0 = config.get @qwen3_moe.attention.query_size : index + %key_value_size0 = config.get @qwen3_moe.attention.key_value_size : index + %head_size0 = config.get @qwen3_moe.attention.head_size : index + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %query_size, %key_value_size, %head_size = index.assume %query_size0, %key_value_size0, %head_size0 [range(%query_size0, 128, 262144), range(%key_value_size0, 128, 262144), range(%head_size0, 128, 128), mul(%query_size0, %head_size0), mul(%key_value_size0, %head_size0)] : index, index, index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %channel_tile = kernel.workgroup.id : index + %workitem = kernel.workitem.id : index + %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index + %head_tile_count = index.div %head_size, %c8 : index + %head_domain0 = index.div %channel_tile, %head_tile_count : index + %query_head_count = index.div %query_size, %head_size : index + %key_value_head_count = index.div %key_value_size, %head_size : index + %key_value_domain_count = index.mul %key_value_head_count, %c2 : index + %head_domain_count = index.add %query_head_count, %key_value_domain_count : index + %head_domain, %completion_counter_count = index.assume %head_domain0, %head_domain_count [lt(%head_domain0, %head_domain_count)] : index, index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %completion_counters_aligned = buffer.assume.alignment %completion_counters {minimum_alignment = 16} : buffer + %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%completion_counter_count]xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's projection stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete head. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counters_view[%head_domain] {ordering = acq_rel, scope = device} : i32, view<[%completion_counter_count]xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %head_tile_count_i32 = index.cast %head_tile_count : index to i32 + %last_head_tile_i32 = scalar.subi %head_tile_count_i32, %c1_i32 : i32 + %negative_head_tile_count_i32 = scalar.subi %c0_i32, %head_tile_count_i32 : i32 + %is_last_head_tile = scalar.cmpi eq, %old_counter, %last_head_tile_i32 : i32 + scf.if %is_last_head_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_postprocess_output = scalar.constant true : i1 + template.apply<@qwen3_moe.attention.postprocess.head_body>(%publish_postprocess_output, %bounded_token_count, %cache_row_count, %head_domain, %token, %positions, %key_cache_indices, %value_cache_indices, %query_output_raw, %key_output_raw, %value_output_raw, %query_norm_weight, %key_norm_weight, %inverse_frequencies, %query_output, %key_cache, %value_cache) : (i1, index, index, index, index, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer, buffer) + // The counter becomes reusable only after all semantic output stores. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_head_tile_count_i32, %completion_counters_view[%head_domain] {ordering = release, scope = device} : i32, view<[%completion_counter_count]xi32> + } + } + template.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/dense_linear_q4k_next_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/dense_linear_q4k_next_q8.loom new file mode 100644 index 000000000000..30368a61bf05 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/dense_linear_q4k_next_q8.loom @@ -0,0 +1,68 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +template.decl @qwen3_moe.dense.q4k_q8_1_x4.body(%publish_output: i1, %token_count: index, %token0: index, %q8_input: buffer, %weight: buffer, %output: buffer) + +template.decl @qwen3_moe.rmsnorm_quantize_q8_1_x4.body(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer, %arg7: buffer) + +target.decl @qwen3_moe_attention_prepare_gfx11_wave32 + +config.decl @qwen3_moe.dense_quantized.output_size : %value: index where [range(%value, 1, 262144)] + +// Decode-only route that publishes the normalized F32 and Q8_1 x4 rows +// consumed by the following feed-forward boundary. +kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8(%token_count: index) { + %output_size = config.get @qwen3_moe.dense_quantized.output_size : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %c256 = index.constant 256 : index + %output_tiles = index.div %output_size, %c8 : index + kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%token_count: index, %q8_input: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %normalized_output: buffer, %completion_counter: buffer, %next_q8_output: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index + %publish_output = scalar.constant true : i1 + %projection_token = kernel.workgroup.id : index + template.apply<@qwen3_moe.dense.q4k_q8_1_x4.body>(%publish_output, %bounded_token_count, %projection_token, %q8_input, %weight, %output) : (i1, index, index, buffer, buffer, buffer) + %output_size0 = config.get @qwen3_moe.dense_quantized.output_size : index + %output_size = index.assume %output_size0 [range(%output_size0, 128, 32768), mul(%output_size0, 128)] : index + %token0 = kernel.workgroup.id : index + %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index + %c0 = index.constant 0 : index + %c8 = index.constant 8 : index + %workitem = kernel.workitem.id : index + %output_tile_count = index.div %output_size, %c8 : index + %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c0_offset = index.constant 0 : offset + %counter_scratch_bytes = index.constant 4 : offset + %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %normalized_output, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer + %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> + // Publish every producer's residual stores before the leader advances one + // workgroup arrival. The last arrival then acquires the complete row. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> + } + kernel.barrier scope(workgroup) ordering(acq_rel) + %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 + %output_tile_count_i32 = index.cast %output_tile_count : index to i32 + %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 + %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 + %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 + scf.if %is_last_output_tile { + kernel.barrier scope(workgroup) ordering(acquire) + %publish_normalized = scalar.constant true : i1 + template.apply<@qwen3_moe.rmsnorm_quantize_q8_1_x4.body>(%publish_normalized, %c8, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %normalized_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) + // Reset only after every normalized F32 and Q8 store completes. + kernel.barrier scope(workgroup) ordering(release) + scf.if %is_arrival_workitem { + view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/expert_table_partition_prefill_512.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/expert_table_partition_prefill_512.loom new file mode 100644 index 000000000000..80434f33a4d0 --- /dev/null +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/motifs/expert_table_partition_prefill_512.loom @@ -0,0 +1,162 @@ +// Copyright 2026 The HRX Authors +// SPDX-License-Identifier: Apache-2.0 + +amdgpu.target @qwen3_moe_expert_table_partition_gfx11_wave32 {subgroup_size = 32} + +kernel.def target(@qwen3_moe_expert_table_partition_gfx11_wave32) @qwen3_moe_build_expert_table_partition_prefill_512(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { + %c1 = index.constant 1 : index + %launch_expert_count = index.constant 128 : index + %workgroup_size = index.constant 256 : index + kernel.launch.config workgroups(%launch_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index +} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer, %partition_table: buffer, %completion_counter: buffer) { + %bounded_token_count = index.assume %token_count [range(%token_count, 512, 512)] : index + %bounded_route_count = index.assume %route_count [range(%route_count, 8, 8)] : index + %bounded_route_stride = index.assume %route_stride [range(%route_stride, 8, 8)] : index + %bounded_expert_count = index.assume %expert_count [range(%expert_count, 128, 128)] : index + %expert0 = kernel.workgroup.id : index + %expert, %launch_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index + %lane0 = kernel.workitem.id : index + %lane = index.assume %lane0 [range(%lane0, 0, 255)] : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c31 = index.constant 31 : index + %c32 = index.constant 32 : index + %workgroup_size = index.constant 256 : index + %c0_i32 = scalar.constant 0 : i32 + %c1_i32 = scalar.constant 1 : i32 + %c7_i32 = scalar.constant 7 : i32 + %c13_i32 = scalar.constant 13 : i32 + %c0_offset = index.constant 0 : offset + %c4_bytes = index.constant 4 : offset + %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index + %assignment_table_byte_base = index.scale %launch_expert_count, %c4_bytes : index, offset -> offset + %route_ids_noalias, %expert_table_noalias, %partition_table_noalias, %completion_counter_noalias = buffer.assume.noalias %route_ids, %expert_table, %partition_table, %completion_counter : buffer, buffer, buffer, buffer + %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer + %route_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> + %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%launch_expert_count]xi32> + %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%launch_expert_count]x[%bounded_token_count]xi32> + %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> + + %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { + %assignment = index.add %block_base, %lane : index + %in_range = index.cmp ult, %assignment, %assignment_count : index + %route_expert_i32 = scf.if %in_range -> (i32) { + %token0 = index.div %assignment, %bounded_route_count : index + %route0 = index.rem %assignment, %bounded_route_count : index + %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index + %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index + %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 + scf.yield %loaded : i32 + } else { + %cn1_i32 = scalar.constant -1 : i32 + scf.yield %cn1_i32 : i32 + } + %route_expert0 = index.cast %route_expert_i32 : i32 to index + %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 127)] : index + %matches = index.cmp eq, %route_expert, %expert : index + %match_i32 = scf.if %matches -> (i32) { + scf.yield %c1_i32 : i32 + } else { + scf.yield %c0_i32 : i32 + } + %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 + %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 + %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 + scf.if %matches { + %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 + %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index + %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 511)] : index + %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index + %assignment_i32 = index.cast %assignment : index to i32 + view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%launch_expert_count]x[%bounded_token_count]xi32> + } + %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 + scf.yield %next_matched_base : i32 + } + + %is_lane_zero = index.cmp eq, %lane, %c0 : index + scf.if %is_lane_zero { + view.store %expert_route_count, %count_view[%expert] : i32, view<[%launch_expert_count]xi32> + } + + // Every count store precedes its workgroup's release. The last arrival + // acquires all preceding releases before any lane loads the complete table. + %local_old_counter = scf.if %is_lane_zero -> (i32) { + %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 + scf.yield %old_counter : i32 + } else { + scf.yield %c0_i32 : i32 + } + %old_counter = kernel.workgroup.reduce %local_old_counter : i32 + %expert_count_i32 = index.cast %launch_expert_count : index to i32 + %last_expert_i32 = scalar.subi %expert_count_i32, %c1_i32 : i32 + %negative_expert_count_i32 = scalar.subi %c0_i32, %expert_count_i32 : i32 + %is_last_expert = scalar.cmpi eq, %old_counter, %last_expert_i32 : i32 + scf.if %is_last_expert { + %rounded_assignment_count = index.add %assignment_count, %c31 : index + %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index + %maximum_partition_count = index.add %assignment_partition_count, %launch_expert_count : index + %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> + %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> + %has_expert = index.cmp ult, %lane, %launch_expert_count : index + %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { + %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index + %loaded = view.load %count_view[%partition_expert] : view<[%launch_expert_count]xi32> -> i32 + scf.yield %loaded : i32 + } else { + scf.yield %c0_i32 : i32 + } + %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index + %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 512)] : index + %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index + %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index + %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 + %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 + %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 + %partition_count = index.cast %partition_count_i32 : i32 to index + %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index + %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index + %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 255)] : index + scf.if %has_expert { + %partition_expert, %table_expert_count = index.assume %lane, %launch_expert_count [lt(%lane, %launch_expert_count)] : index, index + %expert_i32 = index.cast %partition_expert : index to i32 + scf.for %partition = [%c0 to %expert_partition_count step %c1] { + %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index + %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index + %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index + %partition_remainder = index.rem %expert_assignment_count, %c32 : index + %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index + %partition_row_count = scf.if %has_partial_tail -> (index) { + %next_partition = index.add %partition, %c1 : index + %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index + %tail_row_count = scf.if %is_tail_partition -> (index) { + scf.yield %partition_remainder : index + } else { + scf.yield %c32 : index + } + scf.yield %tail_row_count : index + } else { + scf.yield %c32 : index + } + %partition_i32 = index.cast %partition : index to i32 + %partition_row_count_i32 = index.cast %partition_row_count : index to i32 + %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 + %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 + %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 + %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 + %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 + view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> + } + } + scf.if %is_lane_zero { + view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> + } + // The counter cannot become reusable until every descriptor store is + // complete and visible to the following dispatch. + kernel.barrier scope(workgroup) ordering(acq_rel) + scf.if %is_lane_zero { + view.atomic.reduce %negative_expert_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> + } + } + kernel.return +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom deleted file mode 100644 index c853e2c84e28..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_next_q8.loom +++ /dev/null @@ -1,55 +0,0 @@ -// Shared decode completion protocol for routed-down projections that publish -// the normalized Q8_1 x4 row consumed by the next projection boundary. Each -// x workgroup must publish one output tile before applying this body and must -// provide both its tile width and the number of target subgroups participating -// in the row reduction. The last arrival acquires the complete residual row, -// normalizes and packs it, then resets the reusable completion word after every -// Q8 store. -// -// This outer device template deliberately applies the RMSNorm/Q8 device -// template. Keeping that composition authored here ensures the compiler -// preserves the eventual kernel ancestor through nested template selection. -func.template device @qwen3_moe_routed_down_next_q8_completion(%output_channels_per_workgroup0: index, %reduction_subgroup_count0: index, %token_count: index, %output_size: index, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %output_channels_per_workgroup = index.assume %output_channels_per_workgroup0 [range(%output_channels_per_workgroup0, 1, 8)] : index - %reduction_subgroup_count = index.assume %reduction_subgroup_count0 [range(%reduction_subgroup_count0, 1, 8)] : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 128, 32768), mul(%output_size, 128)] : index - %token0 = kernel.workgroup.id : index - %token = index.assume %token0 [lt(%token0, %bounded_token_count)] : index - %c0 = index.constant 0 : index - %workitem = kernel.workitem.id : index - %output_tile_count = index.div %bounded_output_size, %output_channels_per_workgroup : index - %is_arrival_workitem = index.cmp eq, %workitem, %c0 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %output_noalias, %norm_weight_noalias, %completion_counter_noalias, %next_q8_output_noalias = buffer.assume.noalias %output, %norm_weight, %completion_counter, %next_q8_output : buffer, buffer, buffer, buffer - %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer - %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - // Publish every producer's residual stores before the leader advances one - // workgroup arrival. The last arrival then acquires the complete row. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - %old_counter = view.atomic.rmw %c1_i32, %completion_counter_view[%c0] {ordering = acq_rel, scope = device} : i32, view<1xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %output_tile_count_i32 = index.cast %output_tile_count : index to i32 - %last_output_tile_i32 = scalar.subi %output_tile_count_i32, %c1_i32 : i32 - %negative_output_tile_count_i32 = scalar.subi %c0_i32, %output_tile_count_i32 : i32 - %is_last_output_tile = scalar.cmpi eq, %old_counter, %last_output_tile_i32 : i32 - scf.if %is_last_output_tile { - kernel.barrier scope(workgroup) ordering(acquire) - %publish_normalized = scalar.constant false : i1 - func.apply(%publish_normalized, %reduction_subgroup_count, %bounded_token_count, %token, %output_noalias, %norm_weight_noalias, %next_q8_output_noalias, %next_q8_output_noalias) : (i1, index, index, index, buffer, buffer, buffer, buffer) - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_workitem { - view.atomic.reduce %negative_output_tile_count_i32, %completion_counter_view[%c0] {ordering = release, scope = device} : i32, view<1xi32> - } - } - func.return -} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom deleted file mode 100644 index df599d53dd80..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q4k.loom +++ /dev/null @@ -1,395 +0,0 @@ -// Fuses the Qwen3 MoE Q4_K down projection with route weighting, top-8 -// reduction, and residual publication. The input contains one Q8_1 x4 row per -// [token, route], while weights remain in GGUF's raw -// [expert, output, K / 256, 144 bytes] layout. -// -// Route IDs retain an independent physical stride because llama.cpp selects a -// top-8 view from its 128-entry argsort storage. Normalized route weights are -// compact [token, route]. The output enters containing the residual and is -// updated in place, so the fused boundary never materializes an unweighted or -// route-indexed [token, route, hidden] down-projection tensor. -// -// One wave owns one output channel and contracts all selected routes in -// registers. Four eight-lane cohorts contract four independent routes while -// walking every Q4_K block. Qwen's eight-route, three-block decode shape uses -// all 32 lanes across two route batches without addressing synthetic blocks. -target.decl @qwen3_moe_attention_prepare_gfx11_wave32 - -config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) - -kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %q8_output: buffer) - -// Shared Q4_K contraction and residual publication. Each lane consumes both -// nibbles of a packed code load before the subgroup reduces the row. -func.template device @qwen3_moe_routed_down_q4k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144 = index.constant 144 : index - %c256 = index.constant 256 : index - %c1_byte = index.constant 1 : offset - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %cohort0 = index.div %lane, %c8 : index - %cohort = index.assume %cohort0 [range(%cohort0, 0, 3)] : index - %cohort_lane0 = index.rem %lane, %c8 : index - %cohort_lane = index.assume %cohort_lane0 [range(%cohort_lane0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_byte_count = index.mul %q4_block_count, %c144 : index - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_byte_count = index.mul %q8_group_count, %c144 : index - %padded_route_count = index.add %bounded_route_count, %c3 : index - %route_batch_count = index.div %padded_route_count, %c4 : index - %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> - %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index - %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index - %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 - %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 - %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { - %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index - %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - scf.yield %expert_i32, %route_weight : i32, f32 - } else { - scf.yield %c0_i32, %c0_f32 : i32, f32 - } - %computes_channel = scalar.andi %publish_output, %valid_channel : i1 - %routed_lane_sum = scf.if %computes_channel -> (f32) { - %sum = scf.for %route_batch = [%c0 to %route_batch_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { - %route_batch_base = index.mul %route_batch, %c4 : index - %route0 = index.add %route_batch_base, %cohort : index - %route = index.assume %route0 [range(%route0, 0, 7)] : index - %active_route = index.cmp ult, %route, %bounded_route_count : index - %weighted_lane = scf.if %active_route -> (f32) { - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index - %route_i32 = index.cast %bounded_route : index to i32 - %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %bounded_output_size : index - %expert_channel = index.add %expert_output_base, %channel : index - %weight_row_byte_index = index.mul %expert_channel, %weight_row_byte_count : index - %weight_row_byte_base = index.scale %weight_row_byte_index, %c1_byte : index, offset -> offset - %q8_row_base0 = index.mul %bounded_token, %body_route_count : index - %q8_row = index.add %q8_row_base0, %bounded_route : index - %q8_row_byte_index = index.mul %q8_row, %q8_row_byte_count : index - %q8_row_byte_base = index.scale %q8_row_byte_index, %c1_byte : index, offset -> offset - %route_lane_sum = func.call @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %cohort_lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 - %weighted = scalar.mulf %route_lane_sum, %route_weight : f32 - scf.yield %weighted : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %route_acc, %weighted_lane : f32 - scf.yield %next : f32 - } - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 - %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 - %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 - scf.if %writes_output { - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 - %result = scalar.addf %residual, %routed_sum : f32 - view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> - } - func.return -} - -kernel.def @qwen3_moe_routed_down_q4k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %c0 = index.constant 0 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - func.apply(%valid_token, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Decode-only route that also publishes the normalized Q8_1 x4 row consumed by -// the next projection boundary. -kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index - %c8 = index.constant 8 : index - %publishes_output = scalar.constant true : i1 - %body_token = index.constant 0 : index - func.apply(%publishes_output, %bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - func.apply(%c8, %c8, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// The production hidden width requires all 256 output tiles to arrive before -// normalization. Compare both residual and packed Q8 publication with the -// ordinary two-dispatch composition, then reuse the same completion word. -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(2048) : index - %routed_row_count = check.literal value(2) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.iota offset(-72) step(1) period(144) : tensor<2x2048x3x144xi8> - %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> - %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.000001) rtol(0.000001) nan(same) : tensor<1x2048xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.000001) rtol(0.000001) nan(same) : tensor<1x2048xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x144xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x144xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %next_q8_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - check.return -} - -// Two selected experts exercise the noncompact route-ID stride, normalized -// weighted reduction, in-place residual update, the odd Q4_K block tail at -// K=768, and the output tile tail. Uniform 0xaa bytes decode to a deterministic -// nonzero row. -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_residual_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(9) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x9x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<1x9xf32> - %expected = check.generate.fill value(-58.0394287109375) : tensor<1x9xf32> - %routed_row_count = check.literal value(2) : index - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x144xi8>, tensor<1x9xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<1x9xf32> - check.return -} - -// Two tokens use different route-weight sums so token, route, packed-row, and -// output addressing cannot collapse to the single-token case. -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_prefill_case { - %token_count = check.literal value(2) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(1) : index - %routed_row_count = check.literal value(4) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<2x4xi32> - %route_weights = check.generate.iota offset(0.25) step(0.25) period(4) : tensor<2x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x1x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<2x1xf32> - %expected = check.generate.iota offset(-43.279571533203125) step(-59.0394287109375) : tensor<2x1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x2x768xf32>, tensor<2x2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x2x864xi8>, tensor<2x4xi32>, tensor<2x2xf32>, tensor<2x1x3x144xi8>, tensor<2x1xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x1xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - // Each 768-element route contains six complete Q8_1 x4 groups, so packing - // the eight contiguous routes as one 6,144-element token row preserves the - // exact per-route physical layout without a derived testbench scalar. - %routed_input_size = check.literal value(6144) : index - %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %routed_input_size](%token_count, %routed_input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x864xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_nonzero_residual_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_pipeline_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_decode - -check.benchmark<@qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_down_q4k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom deleted file mode 100644 index e59e3e07ff1a..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_q6k.loom +++ /dev/null @@ -1,737 +0,0 @@ -// Fuses the Qwen3 MoE Q6_K down projection with route weighting, top-8 -// reduction, and residual publication. The input contains one Q8_1 x4 row per -// [token, route], while weights remain in GGUF's raw -// [expert, output, K / 256, 210 bytes] layout. -// -// Route IDs retain an independent physical stride because llama.cpp selects a -// top-8 view from its 128-entry argsort storage. Normalized route weights are -// compact [token, route]. The output enters containing the residual and is -// updated in place, so the fused boundary never materializes an unweighted or -// route-indexed [token, route, hidden] down-projection tensor. -// -// The Q8_1 provider assigns one output channel to a wave and reduces all routes -// in registers. The direct-F32 provider mirrors llama.cpp's gfx1151 Vulkan -// decode algorithm: four wave64 subgroups compute four channels, each subgroup -// contracts four raw Q6_K blocks at a time, and route weighting and residual -// publication remain fused. A grouped prefill provider can reuse the same raw -// Q6_K row primitive while staging each expert row across multiple tokens. -amdgpu.target @qwen3_moe_routed_down_q6k_gfx11_wave64 {subgroup_size = 64} - -target.decl @qwen3_moe_attention_prepare_gfx11_wave32 - -config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] - -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -func.decl @ggml_q6k_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_q6k_stage_f32_scales(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer) - -func.decl @ggml_q6k_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>) - -func.decl @ggml_q6k_f32_block_row(%input_size: index, %row: index, %block: index, %frame_count: index, %frame: index, %lane: index, %weight: buffer, %scale_stage: buffer, %input0: vector<4xf32>, %input1: vector<4xf32>, %input2: vector<4xf32>, %input3: vector<4xf32>) -> (f32) - -kernel.decl @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %q8_output: buffer) - -// Shared Q6_K contraction and residual publication used by ordinary and -// completion-fused exports with identical launch geometry. -func.template device @qwen3_moe_routed_down_q6k_q8_1_x4_body(%publish_output: i1, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %c144_bytes = index.constant 144 : offset - %c210_bytes = index.constant 210 : offset - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %channel_base = index.mul %channel_tile, %c8 : index - %channel = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %q6_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.scale %q6_block_count, %c210_bytes : index, offset -> offset - %q8_group_count = index.div %bounded_input_size, %c128 : index - %q8_row_bytes = index.scale %q8_group_count, %c144_bytes : index, offset -> offset - %q8_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> - %lane_has_route = index.cmp ult, %lane, %bounded_route_count : index - %lane_within_route_stride = index.cmp ult, %lane, %bounded_route_id_stride : index - %loads_active_route = scalar.andi %publish_output, %lane_has_route : i1 - %loads_route_metadata = scalar.andi %loads_active_route, %lane_within_route_stride : i1 - %lane_expert_i32, %lane_route_weight = scf.if %loads_route_metadata -> (i32, f32) { - %route_lane0, %metadata_route_count = index.assume %lane, %bounded_route_count [lt(%lane, %bounded_route_count)] : index, index - %route_lane, %metadata_route_id_stride = index.assume %route_lane0, %bounded_route_id_stride [lt(%route_lane0, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_lane] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %route_lane0] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - scf.yield %expert_i32, %route_weight : i32, f32 - } else { - scf.yield %c0_i32, %c0_f32 : i32, f32 - } - %computes_channel = scalar.andi %publish_output, %valid_channel : i1 - %routed_lane_sum = scf.if %computes_channel -> (f32) { - %sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { - %route_i32 = index.cast %route : index to i32 - %expert_i32 = kernel.subgroup.broadcast %lane_expert_i32 from %route_i32 : i32, i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %bounded_output_size : index - %expert_channel = index.add %expert_output_base, %channel : index - %weight_row_byte_base = index.scale %expert_channel, %weight_row_bytes : index, offset -> offset - %q8_row_base0 = index.mul %bounded_token, %bounded_route_count : index - %q8_row = index.add %q8_row_base0, %route : index - %q8_row_byte_base = index.scale %q8_row, %q8_row_bytes : index, offset -> offset - %lane_acc = func.call @ggml_q6k_q8_1_x4_row_lane(%bounded_input_size, %weight_noalias, %weight_row_byte_base, %q8_noalias, %q8_row_byte_base, %lane) : (index, buffer, offset, buffer, offset, index) -> (f32) - %route_weight = kernel.subgroup.broadcast %lane_route_weight from %route_i32 : f32, i32 - %weighted_lane = scalar.mulf %lane_acc, %route_weight : f32 - %next = scalar.addf %route_acc, %weighted_lane : f32 - scf.yield %next : f32 - } - scf.yield %sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %routed_sum = kernel.subgroup.reduce %routed_lane_sum : f32 - %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 - %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 - scf.if %writes_output { - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 - %result = scalar.addf %residual, %routed_sum : f32 - view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> - } - func.return -} - -kernel.def @qwen3_moe_routed_down_q6k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %c0 = index.constant 0 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - func.apply(%valid_token, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Decode-only route that also publishes the normalized Q8_1 x4 row consumed by -// the next projection boundary. -kernel.def target(@qwen3_moe_attention_prepare_gfx11_wave32) @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c8 = index.constant 8 : index - %c7 = index.constant 7 : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index - %c8 = index.constant 8 : index - %publishes_output = scalar.constant true : i1 - %body_token = index.constant 0 : index - func.apply(%publishes_output, %bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - func.apply(%c8, %c8, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// The production hidden width requires all 256 output tiles to arrive before -// normalization. Compare both residual and packed Q8 publication with the -// ordinary two-dispatch composition, then reuse the same completion word. -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(2048) : index - %routed_row_count = check.literal value(2) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x2048x3x210xi8> - %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> - %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<1x2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %expected_output, %norm_weight, %expected_q8) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output, %norm_weight, %completion_counter, %next_q8_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x210xi8>, tensor<1x2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_benchmark_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<1x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<1x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<1x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<1x2048xf32> - %norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %next_q8_output = check.generate.fill value(0) : tensor<2304xi8> - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<1x8x864xi8>, tensor<1x8xi32>, tensor<1x8xf32>, tensor<128x2048x3x210xi8>, tensor<1x2048xf32>) - kernel.launch @qwen3_moe_attention_rmsnorm_quantize_q8_1_x4[%token_count](%token_count, %output, %norm_weight, %next_q8_output) : [index](index, tensor<1x2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - check.return -} - -// Shared direct-F32 decode schedule. Every subgroup owns one scale frame and -// one output channel. Providers specialize how many Q6_K blocks a subgroup -// covers per pass from the target subgroup width. -func.template device @qwen3_moe_routed_down_q6k_f32_body(%publish_output: i1, %subgroup_count: index, %block_step: index, %scale_stage_bytes: offset, %token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 262144)] : index - %channel_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane0 = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c16 = index.constant 16 : index - %c210 = index.constant 210 : index - %c256 = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index - %channel_tile_base = index.mul %channel_tile, %subgroup_count : index - %channel = index.add %channel_tile_base, %subgroup : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %safe_channel = scf.select %valid_channel, %channel, %c0 : index - %lane_i32 = index.cast %lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %assignment_count = index.mul %body_token_count, %bounded_route_count : index - %q6_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_byte_count = index.mul %q6_block_count, %c210 : index - %cohort = index.div %lane, %c16 : index - %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_output_size]xf32> - %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer - %routed_sum0 = scf.if %publish_output -> (f32) { - %active_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc = %c0_f32 : f32) -> (f32) unroll { - // Route metadata is uniform across all channels. Keeping the access in - // the unrolled route body exposes that uniformity directly to target - // lowering instead of materializing a dynamic subgroup broadcast. - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index - %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %weight_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %bounded_output_size : index - %expert_channel = index.add %expert_output_base, %safe_channel : index - %input_row_base = index.mul %bounded_token, %bounded_route_count : index - %input_row = index.add %input_row_base, %bounded_route : index - %lane_sum = scf.for %block_base = [%c0 to %q6_block_count step %block_step](%block_acc = %c0_f32 : f32) -> (f32) { - %block = index.add %block_base, %cohort : index - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_channel, %block, %subgroup_count, %subgroup, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - %input0, %input1, %input2, %input3 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_channel, %block, %subgroup_count, %subgroup, %lane, %weight_noalias, %scale_stage, %input0, %input1, %input2, %input3) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - %next = scalar.addf %block_acc, %contribution : f32 - scf.yield %next : f32 - } - %route_sum = kernel.subgroup.reduce %lane_sum : f32 - %weighted = scalar.mulf %route_sum, %route_weight : f32 - %next = scalar.addf %route_acc, %weighted : f32 - scf.yield %next : f32 - } - scf.yield %active_sum : f32 - } else { - scf.yield %c0_f32 : f32 - } - %routed_sum = scf.select %valid_channel, %routed_sum0, %c0_f32 : f32 - %writes_active_channel = scalar.andi %publish_output, %valid_channel : i1 - %writes_output = scalar.andi %writes_active_channel, %is_lane_zero : i1 - scf.if %writes_output { - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %residual = view.load %output_view[%bounded_token, %bounded_channel] : view<[%body_token_count]x[%bounded_output_size]xf32> -> f32 - %result = scalar.addf %residual, %routed_sum : f32 - view.store %result, %output_view[%bounded_token, %bounded_channel] : f32, view<[%body_token_count]x[%bounded_output_size]xf32> - } - func.return -} - -// Decode schedule matching the Vulkan oracle's two adjacent output rows per -// wave. Four wave64 subgroups publish eight channels per workgroup. Keeping the -// route loop rolled prevents its eight iterations from multiplying the two-row -// contraction's register pressure and instruction footprint. -func.template device @qwen3_moe_routed_down_q6k_f32_pair_body(%token_count: index, %token: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 512)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 2, 4096)] : index - %pair_tile = kernel.workgroup.id : index - %subgroup0 = kernel.subgroup.id : index - %lane0 = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %c0_offset = index.constant 0 : offset - %scale_stage_bytes = index.constant 2048 : offset - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 3)] : index - %lane = index.assume %lane0 [range(%lane0, 0, 63)] : index - %pair_tile_base = index.mul %pair_tile, %c4 : index - %pair = index.add %pair_tile_base, %subgroup : index - %row00 = index.mul %pair, %c2 : index - %row0, %body_output_size = index.assume %row00, %bounded_output_size [lt(%row00, %bounded_output_size)] : index, index - %row1 = index.add %row0, %c1 : index - %row1_valid = index.cmp ult, %row1, %body_output_size : index - %safe_row1 = scf.select %row1_valid, %row1, %c0 : index - %assignment_count = index.mul %body_token_count, %bounded_route_count : index - %q6_block_count = index.div %bounded_input_size, %c256 : index - %cohort = index.div %lane, %c16 : index - %scale_frame0 = index.mul %subgroup, %c2 : index - %scale_frame1 = index.add %scale_frame0, %c1 : index - %input_noalias, %route_id_noalias, %route_weight_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %route_ids, %route_weights, %weight, %output : buffer, buffer, buffer, buffer, buffer - %route_id_view = buffer.view %route_id_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_id_stride]xi32> - %route_weight_view = buffer.view %route_weight_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%bounded_route_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%body_token_count]x[%body_output_size]xf32> - %scale_stage = buffer.alloca align(16) %scale_stage_bytes : buffer - %routed_sum0, %routed_sum1 = scf.for %route = [%c0 to %bounded_route_count step %c1](%route_acc0 = %c0_f32 : f32, %route_acc1 = %c0_f32 : f32) -> (f32, f32) { - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, index - %route_for_stride, %body_route_id_stride = index.assume %bounded_route, %bounded_route_id_stride [lt(%bounded_route, %bounded_route_id_stride)] : index, index - %expert_i32 = view.load %route_id_view[%bounded_token, %route_for_stride] : view<[%body_token_count]x[%bounded_route_id_stride]xi32> -> i32 - %route_weight = view.load %route_weight_view[%bounded_token, %bounded_route] : view<[%body_token_count]x[%bounded_route_count]xf32> -> f32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert1 = index.assume %expert0 [range(%expert0, 0, 511)] : index - %expert, %body_expert_count = index.assume %expert1, %bounded_expert_count [lt(%expert1, %bounded_expert_count)] : index, index - %expert_output_base = index.mul %expert, %body_output_size : index - %expert_row0 = index.add %expert_output_base, %row0 : index - %expert_row1 = index.add %expert_output_base, %safe_row1 : index - %input_row_base = index.mul %bounded_token, %body_route_count : index - %input_row = index.add %input_row_base, %bounded_route : index - %lane_sum0, %lane_sum1 = scf.for %block_base = [%c0 to %q6_block_count step %c4](%block_acc0 = %c0_f32 : f32, %block_acc1 = %c0_f32 : f32) -> (f32, f32) { - %block = index.add %block_base, %cohort : index - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - %input00, %input01, %input02, %input03 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %contribution0 = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row0, %block, %c8, %scale_frame0, %lane, %weight_noalias, %scale_stage, %input00, %input01, %input02, %input03) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - %contribution1 = scf.if %row1_valid -> (f32) { - func.call @ggml_q6k_stage_f32_scales(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage) : (index, index, index, index, index, index, buffer, buffer) - // Reloading the shared activations lets the first row's vectors die - // before the second scale stage. Carrying them across that stage - // increases register pressure and is slower on gfx1151. - %input10, %input11, %input12, %input13 = func.call @ggml_q6k_load_f32_block(%assignment_count, %bounded_input_size, %input_row, %block, %lane, %input_noalias) : (index, index, index, index, index, buffer) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) - %row1_contribution = func.call @ggml_q6k_f32_block_row(%bounded_input_size, %expert_row1, %block, %c8, %scale_frame1, %lane, %weight_noalias, %scale_stage, %input10, %input11, %input12, %input13) : (index, index, index, index, index, index, buffer, buffer, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) -> (f32) - scf.yield %row1_contribution : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next0 = scalar.addf %block_acc0, %contribution0 : f32 - %next1 = scalar.addf %block_acc1, %contribution1 : f32 - scf.yield %next0, %next1 : f32, f32 - } - %route_sum0 = kernel.subgroup.reduce %lane_sum0 : f32 - %route_sum1 = kernel.subgroup.reduce %lane_sum1 : f32 - %weighted0 = scalar.mulf %route_sum0, %route_weight : f32 - %weighted1 = scalar.mulf %route_sum1, %route_weight : f32 - %next0 = scalar.addf %route_acc0, %weighted0 : f32 - %next1 = scalar.addf %route_acc1, %weighted1 : f32 - scf.yield %next0, %next1 : f32, f32 - } - %is_lane_zero = index.cmp eq, %lane, %c0 : index - scf.if %is_lane_zero { - %residual0 = view.load %output_view[%bounded_token, %row0] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 - %result0 = scalar.addf %residual0, %routed_sum0 : f32 - view.store %result0, %output_view[%bounded_token, %row0] : f32, view<[%body_token_count]x[%body_output_size]xf32> - scf.if %row1_valid { - %residual1 = view.load %output_view[%bounded_token, %row1] : view<[%body_token_count]x[%body_output_size]xf32> -> f32 - %result1 = scalar.addf %residual1, %routed_sum1 : f32 - view.store %result1, %output_view[%bounded_token, %row1] : f32, view<[%body_token_count]x[%body_output_size]xf32> - } - } - func.return -} - -// Wave64 matches llama.cpp's Vulkan subgroup width and contracts four Q6_K -// blocks per subgroup pass. -kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_routed_down_q6k_f32_wave64(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c3 : index - %output_tiles = index.div %padded_output_size, %c4 : index - kernel.launch.config workgroups(%output_tiles, %token_capacity, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %c0 = index.constant 0 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %c4 = index.constant 4 : index - %scale_stage_bytes = index.constant 1024 : offset - func.apply(%valid_token, %c4, %c4, %scale_stage_bytes, %body_token_count, %safe_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %input, %route_ids, %route_weights, %weight, %output) : (i1, index, index, offset, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Decode-only direct-F32 route that publishes the normalized Q8_1 x4 row -// consumed by the next projection boundary. Each wave64 subgroup owns two -// adjacent output channels, so four subgroups publish eight channels while the -// completion epilogue still reduces four subgroup sums. -kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_routed_down_q6k_f32_wave64_next_q8(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) { - %configured_output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c7 = index.constant 7 : index - %c8 = index.constant 8 : index - %c256 = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c7 : index - %output_tiles = index.div %padded_output_size, %c8 : index - kernel.launch.config workgroups(%output_tiles, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer, %norm_weight: buffer, %completion_counter: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %configured_input_size0 = config.get @qwen3_moe.routed_down.input_size : index - %configured_route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_down.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %bounded_input_size, %configured_input_size = index.assume %input_size, %configured_input_size0 [range(%input_size, 256, 32768), mul(%input_size, 256), eq(%input_size, %configured_input_size0)] : index, index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_id_stride = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %body_token = index.constant 0 : index - func.apply(%bounded_token_count, %body_token, %configured_input_size, %configured_route_count, %bounded_route_id_stride, %configured_expert_count, %configured_output_size, %input, %route_ids, %route_weights, %weight, %output) : (index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - func.apply(%c8, %c4, %bounded_token_count, %configured_output_size, %output, %norm_weight, %completion_counter, %next_q8_output) : (index, index, index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// Two selected experts exercise the noncompact route-ID stride, normalized -// weighted reduction, in-place residual update, K=768, and the output tile -// tail. Uniform 0xaa Q6_K bytes decode to a deterministic nonzero row with -// negative signed group scales. -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_residual_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(9) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x9x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<1x9xf32> - %expected = check.generate.fill value(135.31431579589844) : tensor<1x9xf32> - %routed_row_count = check.literal value(2) : index - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<1x9xf32> - check.return -} - -// Two tokens use different route-weight sums so token, route, packed-row, and -// output addressing cannot collapse to the single-token case. -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_prefill_case { - %token_count = check.literal value(2) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(1) : index - %routed_row_count = check.literal value(4) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<2x4xi32> - %route_weights = check.generate.iota offset(0.25) step(0.25) period(4) : tensor<2x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x1x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<2x1xf32> - %expected = check.generate.iota offset(101.73573684692383) step(134.31431579589844) : tensor<2x1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x2x768xf32>, tensor<2x2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x2x864xi8>, tensor<2x4xi32>, tensor<2x2xf32>, tensor<2x1x3x210xi8>, tensor<2x1xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<2x1xf32> - check.return -} - -// The exact-representable routed inputs make direct-F32 and Q8_1 contraction -// comparable while preserving noncompact route IDs and an output-channel tail. -check.case public @qwen3_moe_routed_down_q6k_f32_wave64_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(2) : index - %route_id_stride = check.literal value(4) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(9) : index - %routed_row_count = check.literal value(2) : index - %routed_input = check.generate.fill value(0.00390625) : tensor<2x768xf32> - %q8_input = check.generate.fill value(0) : tensor<2x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(2) : tensor<1x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<1x2xf32> - %weight = check.generate.fill value(-86) : tensor<2x9x3x210xi8> - %expected = check.generate.fill value(1.0) : tensor<1x9xf32> - %actual_wave64 = check.generate.fill value(1.0) : tensor<1x9xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<2x768xf32>, tensor<2x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x864xi8>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %routed_input, %route_ids, %route_weights, %weight, %actual_wave64) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<2x768xf32>, tensor<1x4xi32>, tensor<1x2xf32>, tensor<2x9x3x210xi8>, tensor<1x9xf32>) - check.expect.close actual(%actual_wave64) expected(%expected) atol(0.25) rtol(0.01) nan(same) : tensor<1x9xf32> - check.return -} - -// Wave64 reference for the row producer used by the fused completion path. -kernel.def target(@qwen3_moe_routed_down_q6k_gfx11_wave64) @qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check() { - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%input: buffer, %norm_weight: buffer, %q8_output: buffer) { - %publish_normalized = scalar.constant false : i1 - %reduction_subgroup_count = index.constant 4 : index - %token_count = index.constant 1 : index - %token = index.constant 0 : index - func.apply(%publish_normalized, %reduction_subgroup_count, %token_count, %token, %input, %norm_weight, %q8_output, %q8_output) : (i1, index, index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// The production hidden width requires every direct-F32 output tile to arrive -// before normalization. Nonzero Q6_K data compares the paired-row producer -// against the ordinary wave64 producer and standalone normalization, then -// reuses the completion word so stale completion state is observable. -check.case public @qwen3_moe_routed_down_q6k_f32_wave64_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(8) : index - %expert_count = check.literal value(8) : index - %output_size = check.literal value(2048) : index - %input = check.generate.iota offset(-0.5) step(0.00016276041666666666) period(6144) : tensor<8x768xf32> - %route_ids = check.generate.iota offset(0) step(1) period(8) : tensor<8xi32> - %route_weights = check.generate.iota offset(0.027777777777777776) step(0.027777777777777776) : tensor<8xf32> - %weight = check.generate.fill value(-86) : tensor<8x2048x3x210xi8> - %norm_weight = check.generate.iota offset(-1.0) step(0.0009765625) : tensor<2048xf32> - %expected_output = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> - %actual_output0 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> - %actual_output1 = check.generate.iota offset(-0.5) step(0.00048828125) period(2048) : tensor<2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<2304xi8> - %actual_q8_0 = check.generate.fill value(1) : tensor<2304xi8> - %actual_q8_1 = check.generate.fill value(1) : tensor<2304xi8> - %completion_counter = check.generate.fill value(0) : tensor<1xi32> - %expected_counter = check.generate.fill value(0) : tensor<1xi32> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>) - kernel.launch @qwen3_moe_rmsnorm_quantize_q8_1_x4_wave64_production_check(%expected_output, %norm_weight, %expected_q8) : (tensor<2048xf32>, tensor<2048xf32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %actual_output0, %norm_weight, %completion_counter, %actual_q8_0) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64_next_q8[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %input, %route_ids, %route_weights, %weight, %actual_output1, %norm_weight, %completion_counter, %actual_q8_1) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<8x768xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8x2048x3x210xi8>, tensor<2048xf32>, tensor<2048xf32>, tensor<1xi32>, tensor<2304xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<2048xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<2048xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<2304xi8> - check.expect.equal actual(%completion_counter) expected(%expected_counter) : tensor<1xi32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - // Each 768-element route contains six complete Q8_1 x4 groups, so packing - // the eight contiguous routes as one 6,144-element token row preserves the - // exact per-route physical layout without a derived testbench scalar. - %routed_input_size = check.literal value(6144) : index - %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x8x864xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %routed_input_size](%token_count, %routed_input_size, %routed_input, %q8_input) : [index, index](index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x864xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x864xi8>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(768) : index - %route_count = check.literal value(8) : index - %route_id_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(2048) : index - %routed_input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_routed_down_q6k_f32_wave64[%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_id_stride, %expert_count, %output_size, %routed_input, %route_ids, %route_weights, %weight, %output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x128xi32>, tensor<[%token_count]x8xf32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_nonzero_residual_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_pipeline_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_f32_wave64_benchmark_case> @qwen3_moe_routed_down_q6k_f32_wave64_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_decode - -check.benchmark<@qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_down_q6k_q8_1_x4_next_q8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom index 84dcb34bd94b..4b1a879bde48 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_quantized_f16_wmma.loom @@ -1,553 +1,16 @@ -// Expert-grouped Q4_K and Q6_K down projections for gfx11 prefill. +// Weighted reduction of compact FP16 routed projections into an FP32 residual. // -// Each two-wave workgroup decodes 64 output channels for 32 routed rows of one -// expert. Raw GGUF weights and FP16 routed SwiGLU rows are staged as FP16 WMMA -// operands. Results remain FP16 and are scattered into compact -// [token, route, hidden] order without collisions. A following reduction -// widens each route, applies its normalized weight, and accumulates the -// residual in FP32. -// -// Top-k routing selects an expert at most once per token. The expert table -// therefore contains at most token_count assignments per expert, while each -// assignment ordinal still identifies the compact [token, route] activation -// and route-weight row. -amdgpu.target @qwen3_moe_routed_down_gfx11_wave64 {subgroup_size = 64} +// One workitem owns four adjacent output channels. All top-k routes are +// accumulated in FP32 before the residual is loaded and published once. -config.decl @qwen3_moe.routed_down.input_size : %value: index where [range(%value, 256, 32768), mul(%value, 256)] +template.decl @ggml.routing.weighted_reduce_f16_f32(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] +amdgpu.target @qwen3_moe_routed_down_gfx11_wave64 {subgroup_size = 64} -config.decl @qwen3_moe.routed_down.expert_count : %value: index where [range(%value, 1, 512)] +config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) - -kernel.decl @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) - -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) - -func.decl @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) - -func.decl @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) - -func.decl @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) - -func.decl @ggml_q6k_f16_vector4(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf16>) - -kernel.decl @qwen3_moe_routed_down_q4k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -kernel.decl @qwen3_moe_routed_down_q6k_q8_1_x4(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %input_size: index, %route_count: index, %route_id_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %route_weights: buffer, %weight: buffer, %output: buffer) - -// Acquires the packed words shared by one four-group half of a Q6_K block. -// Each QL word supplies two groups and the QH word supplies all four. -func.def inline @qwen3_moe_q6k_wmma_load_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_half: index, %packet: index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 210 : offset - %qh_byte_add = index.constant 128 : offset - %bounded_half = index.assume %q6_half [range(%q6_half, 0, 1)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %qh_byte_base = index.add %block_byte_base, %qh_byte_add : offset - %ql_view = buffer.view %weight[%block_byte_base] : buffer -> view<32xi32> - %qh_view = buffer.view %weight[%qh_byte_base] : buffer -> view<16xi32> - %ql_half_word_base = index.mul %bounded_half, %c16 : index - %ql0_word_index = index.add %ql_half_word_base, %bounded_packet : index - %ql1_word_base = index.add %ql_half_word_base, %c8 : index - %ql1_word_index = index.add %ql1_word_base, %bounded_packet : index - %qh_half_word_base = index.mul %bounded_half, %c8 : index - %qh_word_index = index.add %qh_half_word_base, %bounded_packet : index - %ql0_word = vector.load %ql_view[%ql0_word_index] : view<32xi32> -> vector<1xi32> - %ql1_word = vector.load %ql_view[%ql1_word_index] : view<32xi32> -> vector<1xi32> - %qh_word = vector.load %qh_view[%qh_word_index] : view<16xi32> -> vector<1xi32> - func.return %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> -} - -// Decodes four adjacent values after the surrounding schedule has selected -// the scale and retained the packed code words at their natural lifetimes. -func.def inline @qwen3_moe_q6k_wmma_vector4_from_scale_codes(%q6_group: index, %scale_i8: i8, %d_f16: f16, %ql_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c4_i32v = vector.constant 4 : vector<1xi32> - %nibble_mask = vector.constant 252645135 : vector<1xi32> - %high_mask = vector.constant 50529027 : vector<1xi32> - %c32_f32v = vector.constant 32.0 : vector<4xf32> - %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index - %group_in_half = index.rem %bounded_group, %c4 : index - %nibble = index.div %group_in_half, %c2 : index - %nibble_shift_index = index.mul %nibble, %c4 : index - %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 - %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> - %qh_shift_index = index.mul %group_in_half, %c2 : index - %qh_shift_i32 = index.cast %qh_shift_index : index to i32 - %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> - %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> - %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> - %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> - %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> - %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> - %code = vector.ori %ql, %qh : vector<1xi32> - %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> - %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> - %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> - %scale = scalar.sitofp %scale_i8 : i8 to f32 - %d = scalar.extf %d_f16 : f16 to f32 - %combined_scale = scalar.mulf %scale, %d : f32 - %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> - %values_f32 = vector.mulf %centered, %combined_scale_vector : vector<4xf32> - %values = vector.fptrunc %values_f32 : vector<4xf32> to vector<4xf16> - func.return %values : vector<4xf16> -} - -// Loads only the group scale and block multiplier while reusing packed words -// retained by an enclosing four-group half-block schedule. -func.def inline @qwen3_moe_q6k_wmma_vector4_from_half_codes(%weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index, %ql0_word: vector<1xi32>, %ql1_word: vector<1xi32>, %qh_word: vector<1xi32>) -> (vector<4xf16>) { - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %block_bytes = index.constant 210 : offset - %scale_byte_add = index.constant 192 : offset - %d_byte_add = index.constant 208 : offset - %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q6_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %weight_row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %d_byte_base = index.add %block_byte_base, %d_byte_add : offset - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<16xi8> - %d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16> - %group_in_half = index.rem %bounded_group, %c4 : index - %ql_side = index.rem %group_in_half, %c2 : index - %uses_ql1 = index.cmp eq, %ql_side, %c1 : index - %ql_word = scf.select %uses_ql1, %ql1_word, %ql0_word : vector<1xi32> - %scale_packet_half = index.div %bounded_packet, %c4 : index - %scale_group_base = index.mul %bounded_group, %c2 : index - %scale_index = index.add %scale_group_base, %scale_packet_half : index - %scale_i8 = view.load %scale_view[%scale_index] : view<16xi8> -> i8 - %d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16 - %values = func.call @qwen3_moe_q6k_wmma_vector4_from_scale_codes(%bounded_group, %scale_i8, %d_f16, %ql_word, %qh_word) : (index, i8, f16, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) - func.return %values : vector<4xf16> -} - -// Shared raw-quantized matrix schedule. Entry points pass a literal weight -// format so linking and JIT specialization erase the inactive packed decoder. -func.template device @qwen3_moe_routed_down_quantized_f16_wmma_body(%weight_format: index, %token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.routed_down.input_size : index - %route_count = config.get @qwen3_moe.routed_down.route_count : index - %expert_count = config.get @qwen3_moe.routed_down.expert_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %route_tile = kernel.workgroup.id : index - %expert = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c6 = index.constant 6 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : 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 - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %q6_block_bytes = index.constant 210 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %route_stage_bytes = index.constant 128 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 1024 : offset - %c0_i32 = scalar.constant 0 : i32 - %cn1_i32 = scalar.constant -1 : i32 - %c0_i32x1 = vector.constant 0 : vector<1xi32> - %c0_i32x4 = vector.constant 0 : vector<4xi32> - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %zero_accumulator = vector.constant 0.0 : vector<8xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %assignment_count = index.mul %token_count, %bounded_route_count : index - %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset - %is_q4 = index.cmp eq, %weight_format, %c4 : index - %is_q6 = index.cmp eq, %weight_format, %c6 : index - %quant_block_bytes = scf.select %is_q6, %q6_block_bytes, %q4_block_bytes : offset - %quant_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.scale %quant_block_count, %quant_block_bytes : index, offset -> offset - %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset - %input_noalias, %expert_table_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %weight, %output : buffer, buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%input_size]xf16> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> - %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer - %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer - %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer - %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> - %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> - %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding - %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> - %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> - %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset - %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding - %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> - %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> - %channel_tile_base = index.mul %channel_tile, %c64 : index - %initial_route_tile_base = index.mul %route_tile, %c32 : index - // The launch workload fixes the interleaved route partitions for this exact - // command-program specialization. - %padded_token_count = index.add %token_count, %c63 : index - %route_partition_count = index.div %padded_token_count, %c64 : index - %route_partition_step = index.mul %route_partition_count, %c32 : index - %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index - %is_workitem_zero = index.cmp eq, %workitem, %c0 : index - %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 - %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 - %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index - %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index - scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { - // Snapshot this expert's compact assignment map once, then reuse it for - // every K group and output-channel tile. - %loads_route = index.cmp ult, %workitem, %c32 : index - scf.if %loads_route { - %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index - %assignment_ordinal = index.add %route_tile_base, %local_route : index - %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index - %assignment_i32 = scf.if %valid_row -> (i32) { - %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index - %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %cn1_i32 : i32 - } - view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %load_packet = index.rem %workitem, %c8 : index - %load_k = index.mul %load_packet, %c4 : index - %load_row0 = index.div %workitem, %c8 : index - %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index - %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset - %subgroup_channel_add = index.mul %subgroup, %c32 : index - %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index - %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %result00, %result01, %result10, %result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - // Four lanes cover the 64 output rows owned by one load packet. Snapshot - // each Q4_K block header once and retain it across all eight quant groups. - // The Q6 specialization erases this loop with %is_q4=false. - %header0, %header1, %header2, %header3 = scf.for %header_row_offset = [%c0 to %c64 step %c16](%prior_header0 = %c0_i32x4 : vector<4xi32>, %prior_header1 = %c0_i32x4 : vector<4xi32>, %prior_header2 = %c0_i32x4 : vector<4xi32>, %prior_header3 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { - %header_local_row0 = index.add %load_row, %header_row_offset : index - %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index - %header_channel = index.add %channel_tile_base, %header_local_row : index - %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index - %loads_q4_header = scalar.andi %is_q4, %valid_header_channel : i1 - %loaded_header = scf.if %loads_q4_header -> (vector<4xi32>) { - %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset - %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset - %header_words = func.call @qwen3_moe_q4k_wmma_load_header(%weight_noalias, %header_row_byte_base, %quant_block) : (buffer, offset, index) -> (vector<4xi32>) - scf.yield %header_words : vector<4xi32> - } else { - scf.yield %c0_i32x4 : vector<4xi32> - } - %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index - %updates_header1 = index.cmp eq, %header_row_offset, %c16 : index - %updates_header2 = index.cmp eq, %header_row_offset, %c32 : index - %updates_header3 = index.cmp eq, %header_row_offset, %c48 : index - %next_header0 = scf.select %updates_header0, %loaded_header, %prior_header0 : vector<4xi32> - %next_header1 = scf.select %updates_header1, %loaded_header, %prior_header1 : vector<4xi32> - %next_header2 = scf.select %updates_header2, %loaded_header, %prior_header2 : vector<4xi32> - %next_header3 = scf.select %updates_header3, %loaded_header, %prior_header3 : vector<4xi32> - scf.yield %next_header0, %next_header1, %next_header2, %next_header3 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> - } - %quant_group_outer_count = scf.select %is_q4, %c4, %c2 : index - %groups_per_outer = scf.select %is_q4, %c2, %c4 : index - %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %quant_group_outer = [%c0 to %quant_group_outer_count 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>) { - // Q4_K retains one packed code word across its adjacent low/high group - // pair. The Q6 specialization erases this acquisition path. - %q_word0, %q_word1, %q_word2, %q_word3 = scf.for %code_row_offset = [%c0 to %c64 step %c16](%prior_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_q_word2 = %c0_i32x1 : vector<1xi32>, %prior_q_word3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { - %code_local_row0 = index.add %load_row, %code_row_offset : index - %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index - %code_channel = index.add %channel_tile_base, %code_local_row : index - %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index - %loads_q4_code = scalar.andi %is_q4, %valid_code_channel : i1 - %loaded_q_word = scf.if %loads_q4_code -> (vector<1xi32>) { - %q4_group_pair = index.assume %quant_group_outer [range(%quant_group_outer, 0, 3)] : index - %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset - %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset - %q_word = func.call @qwen3_moe_q4k_wmma_load_code(%weight_noalias, %code_row_byte_base, %quant_block, %q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - scf.yield %q_word : vector<1xi32> - } else { - scf.yield %c0_i32x1 : vector<1xi32> - } - %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index - %updates_q_word1 = index.cmp eq, %code_row_offset, %c16 : index - %updates_q_word2 = index.cmp eq, %code_row_offset, %c32 : index - %updates_q_word3 = index.cmp eq, %code_row_offset, %c48 : index - %next_q_word0 = scf.select %updates_q_word0, %loaded_q_word, %prior_q_word0 : vector<1xi32> - %next_q_word1 = scf.select %updates_q_word1, %loaded_q_word, %prior_q_word1 : vector<1xi32> - %next_q_word2 = scf.select %updates_q_word2, %loaded_q_word, %prior_q_word2 : vector<1xi32> - %next_q_word3 = scf.select %updates_q_word3, %loaded_q_word, %prior_q_word3 : vector<1xi32> - scf.yield %next_q_word0, %next_q_word1, %next_q_word2, %next_q_word3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> - } - // Four Q6_K groups in one half-block share a QH word and consume two - // nibbles from each of two QL words. Retain the three packed words for - // all four groups. The Q4 specialization erases this acquisition path. - %q6_ql00, %q6_ql10, %q6_qh0, %q6_ql01, %q6_ql11, %q6_qh1, %q6_ql02, %q6_ql12, %q6_qh2, %q6_ql03, %q6_ql13, %q6_qh3 = scf.for %q6_code_row_offset = [%c0 to %c64 step %c16](%prior_q6_ql00 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql10 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh0 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql01 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql11 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh1 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql02 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql12 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh2 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql03 = %c0_i32x1 : vector<1xi32>, %prior_q6_ql13 = %c0_i32x1 : vector<1xi32>, %prior_q6_qh3 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { - %q6_code_local_row0 = index.add %load_row, %q6_code_row_offset : index - %q6_code_local_row = index.assume %q6_code_local_row0 [range(%q6_code_local_row0, 0, 63)] : index - %q6_code_channel = index.add %channel_tile_base, %q6_code_local_row : index - %valid_q6_code_channel = index.cmp ult, %q6_code_channel, %bounded_output_size : index - %loads_q6_code = scalar.andi %is_q6, %valid_q6_code_channel : i1 - %loaded_q6_ql0, %loaded_q6_ql1, %loaded_q6_qh = scf.if %loads_q6_code -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) { - %q6_half = index.assume %quant_group_outer [range(%quant_group_outer, 0, 1)] : index - %q6_code_channel_byte_add = index.scale %q6_code_channel, %weight_row_bytes : index, offset -> offset - %q6_code_row_byte_base = index.add %expert_byte_base, %q6_code_channel_byte_add : offset - %ql0_word, %ql1_word, %qh_word = func.call @qwen3_moe_q6k_wmma_load_half_codes(%weight_noalias, %q6_code_row_byte_base, %quant_block, %q6_half, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>) - scf.yield %ql0_word, %ql1_word, %qh_word : vector<1xi32>, vector<1xi32>, vector<1xi32> - } else { - scf.yield %c0_i32x1, %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32>, vector<1xi32> - } - %updates_q6_code0 = index.cmp eq, %q6_code_row_offset, %c0 : index - %updates_q6_code1 = index.cmp eq, %q6_code_row_offset, %c16 : index - %updates_q6_code2 = index.cmp eq, %q6_code_row_offset, %c32 : index - %updates_q6_code3 = index.cmp eq, %q6_code_row_offset, %c48 : index - %next_q6_ql00 = scf.select %updates_q6_code0, %loaded_q6_ql0, %prior_q6_ql00 : vector<1xi32> - %next_q6_ql10 = scf.select %updates_q6_code0, %loaded_q6_ql1, %prior_q6_ql10 : vector<1xi32> - %next_q6_qh0 = scf.select %updates_q6_code0, %loaded_q6_qh, %prior_q6_qh0 : vector<1xi32> - %next_q6_ql01 = scf.select %updates_q6_code1, %loaded_q6_ql0, %prior_q6_ql01 : vector<1xi32> - %next_q6_ql11 = scf.select %updates_q6_code1, %loaded_q6_ql1, %prior_q6_ql11 : vector<1xi32> - %next_q6_qh1 = scf.select %updates_q6_code1, %loaded_q6_qh, %prior_q6_qh1 : vector<1xi32> - %next_q6_ql02 = scf.select %updates_q6_code2, %loaded_q6_ql0, %prior_q6_ql02 : vector<1xi32> - %next_q6_ql12 = scf.select %updates_q6_code2, %loaded_q6_ql1, %prior_q6_ql12 : vector<1xi32> - %next_q6_qh2 = scf.select %updates_q6_code2, %loaded_q6_qh, %prior_q6_qh2 : vector<1xi32> - %next_q6_ql03 = scf.select %updates_q6_code3, %loaded_q6_ql0, %prior_q6_ql03 : vector<1xi32> - %next_q6_ql13 = scf.select %updates_q6_code3, %loaded_q6_ql1, %prior_q6_ql13 : vector<1xi32> - %next_q6_qh3 = scf.select %updates_q6_code3, %loaded_q6_qh, %prior_q6_qh3 : vector<1xi32> - scf.yield %next_q6_ql00, %next_q6_ql10, %next_q6_qh0, %next_q6_ql01, %next_q6_ql11, %next_q6_qh1, %next_q6_ql02, %next_q6_ql12, %next_q6_qh2, %next_q6_ql03, %next_q6_ql13, %next_q6_qh3 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> - } - %outer_result00, %outer_result01, %outer_result10, %outer_result11 = scf.for %group_within_outer = [%c0 to %groups_per_outer step %c1](%group_acc00 = %acc00 : vector<8xf16>, %group_acc01 = %acc01 : vector<8xf16>, %group_acc10 = %acc10 : vector<8xf16>, %group_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %quant_group_base = index.mul %quant_group_outer, %groups_per_outer : index - %quant_group0 = index.add %quant_group_base, %group_within_outer : index - %quant_group = index.assume %quant_group0 [range(%quant_group0, 0, 7)] : index - %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 - %selects_row0 = index.cmp eq, %row_offset, %c0 : index - %selects_row2 = index.cmp eq, %row_offset, %c32 : index - %selects_low_row_pair = index.cmp ult, %row_offset, %c32 : index - %selected_header01 = scf.select %selects_row0, %header0, %header1 : vector<4xi32> - %selected_header23 = scf.select %selects_row2, %header2, %header3 : vector<4xi32> - %selected_header = scf.select %selects_low_row_pair, %selected_header01, %selected_header23 : vector<4xi32> - %selected_q_word01 = scf.select %selects_row0, %q_word0, %q_word1 : vector<1xi32> - %selected_q_word23 = scf.select %selects_row2, %q_word2, %q_word3 : vector<1xi32> - %selected_q_word = scf.select %selects_low_row_pair, %selected_q_word01, %selected_q_word23 : vector<1xi32> - %selected_q6_ql0_01 = scf.select %selects_row0, %q6_ql00, %q6_ql01 : vector<1xi32> - %selected_q6_ql0_23 = scf.select %selects_row2, %q6_ql02, %q6_ql03 : vector<1xi32> - %selected_q6_ql0 = scf.select %selects_low_row_pair, %selected_q6_ql0_01, %selected_q6_ql0_23 : vector<1xi32> - %selected_q6_ql1_01 = scf.select %selects_row0, %q6_ql10, %q6_ql11 : vector<1xi32> - %selected_q6_ql1_23 = scf.select %selects_row2, %q6_ql12, %q6_ql13 : vector<1xi32> - %selected_q6_ql1 = scf.select %selects_low_row_pair, %selected_q6_ql1_01, %selected_q6_ql1_23 : vector<1xi32> - %selected_q6_qh01 = scf.select %selects_row0, %q6_qh0, %q6_qh1 : vector<1xi32> - %selected_q6_qh23 = scf.select %selects_row2, %q6_qh2, %q6_qh3 : vector<1xi32> - %selected_q6_qh = scf.select %selects_low_row_pair, %selected_q6_qh01, %selected_q6_qh23 : vector<1xi32> - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %weight_values = scf.if %valid_channel -> (vector<4xf16>) { - %decoded = scf.if %is_q6 -> (vector<4xf16>) { - %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset - %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset - %q6_values = func.call @qwen3_moe_q6k_wmma_vector4_from_half_codes(%weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet, %selected_q6_ql0, %selected_q6_ql1, %selected_q6_qh) : (buffer, offset, index, index, index, vector<1xi32>, vector<1xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %q6_values : vector<4xf16> - } else { - %q4_values = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%quant_group, %selected_header, %selected_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %q4_values : vector<4xf16> - } - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - %is_activation_row = index.cmp ult, %local_row, %c32 : index - vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - scf.if %is_activation_row { - %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index - %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 - %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 - %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%bounded_assignment, %input_k] : view<[%assignment_count]x[%input_size]xf16> -> vector<4xf16> - scf.yield %loaded : 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 = %group_acc00 : vector<8xf16>, %half_acc01 = %group_acc01 : vector<8xf16>, %half_acc10 = %group_acc10 : vector<8xf16>, %half_acc11 = %group_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 %outer_result00, %outer_result01, %outer_result10, %outer_result11 : 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> - } - // WMMA produces [channel][route] fragments. Transpose through wave-private - // LDS so each lane publishes four adjacent channels for one assignment. - %publish_route0 = index.div %lane, %c4 : index - %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index - %publish_packet0 = index.rem %lane, %c4 : index - %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index - %publish_channel_add = index.mul %publish_packet, %c4 : index - %local_route1 = index.add %c16, %publish_route : index - %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 - %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 - %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 - %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 - %safe_assignment0_i32 = scf.select %assignment0_nonnegative, %assignment0_i32, %c0_i32 : i32 - %safe_assignment1_i32 = scf.select %assignment1_nonnegative, %assignment1_i32, %c0_i32 : i32 - %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index - %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index - %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index - %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index - %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index - %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index - %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index - %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index - %channel1_base = index.add %subgroup_channel_base, %c16 : index - %channel1 = index.add %channel1_base, %publish_channel_add : index - %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index - %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index - %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 - %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 - %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 - %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 - vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes00 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes01 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes10 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes11 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %values, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf16>, view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1> - } - // Route, operand, and result stages are reused by the next concentrated - // routing partition. All waves must finish publication before reuse. - kernel.barrier scope(workgroup) ordering(acq_rel) - } - func.return -} - -// Q4_K and Q6_K retain separate entry points while sharing the complete matrix -// schedule. This keeps format routing outside the hot kernel. -kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_q4k_f16_wmma_grouped(%token_count: index) { - %expert_count = config.get @qwen3_moe.routed_down.expert_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_tiles = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %q4 = index.constant 4 : index - func.apply(%q4, %token_count, %input, %expert_table, %weight, %output) : (index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_q6k_f16_wmma_grouped(%token_count: index) { - %expert_count = config.get @qwen3_moe.routed_down.expert_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_tiles = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %q6 = index.constant 6 : index - func.apply(%q6, %token_count, %input, %expert_table, %weight, %output) : (index, index, buffer, buffer, buffer, buffer) - kernel.return -} - -// Reduces compact routed projections into the residual in logical token order. -// -// One workitem owns four adjacent output channels. All top-k routes are -// accumulated in FP32 before the residual is loaded and published once. The -// preceding grouped projection is the only producer of each FP16 route row, so -// this boundary requires no global atomics. kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) { %output_size = config.get @qwen3_moe.routed_down.output_size : index %c1 = index.constant 1 : index @@ -557,328 +20,24 @@ kernel.def target(@qwen3_moe_routed_down_gfx11_wave64) @qwen3_moe_routed_down_we %rounded_output_size = index.sub %padded_output_size, %c1 : index %output_tiles = index.div %rounded_output_size, %c256 : index kernel.launch.config workgroups(%output_tiles, %token_count, %c1) workgroup_size(%c64, %c1, %c1) : index -} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %route_count = config.get @qwen3_moe.routed_down.route_count : index - %output_size = config.get @qwen3_moe.routed_down.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %token0 = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %assignment_count = index.mul %token_count, %bounded_route_count : index - %channel_tile_base = index.mul %channel_tile, %c256 : index - %lane_channel_add = index.mul %lane, %c4 : index - %channel = index.add %channel_tile_base, %lane_channel_add : index - %active_token = index.cmp ult, %token0, %token_count : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %publishes_output = scalar.andi %active_token, %valid_channel : i1 - %route_weights_noalias, %routed_output_noalias, %output_noalias = buffer.assume.noalias %route_weights, %routed_output, %output : buffer, buffer, buffer - %route_weight_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_route_count]xf32> - %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%bounded_output_size]xf16> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%token_count]x[%bounded_output_size]xf32> - scf.if %publishes_output { - %token, %output_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<4xi1> - %weighted_sum = scf.for %route = [%c0 to %bounded_route_count step %c1](%sum = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll { - %assignment = index.madd %token, %bounded_route_count, %route : index - %routed = vector.load.mask %routed_output_view[%assignment, %channel], %mask, %c0_f16x4 : view<[%assignment_count]x[%bounded_output_size]xf16>, vector<4xi1>, vector<4xf16> - %wide = vector.extf %routed : vector<4xf16> to vector<4xf32> - %route_weight = view.load %route_weight_view[%token, %route] : view<[%token_count]x[%bounded_route_count]xf32> -> f32 - %route_weight_x4 = vector.splat %route_weight : vector<4xf32> - %weighted = vector.mulf %wide, %route_weight_x4 : vector<4xf32> - %next = vector.addf %sum, %weighted : vector<4xf32> - scf.yield %next : vector<4xf32> - } - %residual = vector.load.mask %output_view[%token, %channel], %mask, %c0_f32x4 : view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1>, vector<4xf32> - %result = vector.addf %residual, %weighted_sum : vector<4xf32> - vector.store.mask %result, %output_view[%token, %channel], %mask : vector<4xf32>, view<[%token_count]x[%bounded_output_size]xf32>, vector<4xi1> - } +} launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { + template.apply<@ggml.routing.weighted_reduce_f16_f32>(%token_count, %route_weights, %routed_output, %residual_input, %output) : (index, buffer, buffer, buffer, buffer) kernel.return } -// Complete expert-MLP coverage crosses route partitions and output tails. The -// established Q8_1 gate/up and Q4_K down providers form an independent -// reference while the production path keeps the routed SwiGLU handoff in F16. -// The complete-chain tolerance covers both independent quantized accumulation -// regimes; the focused gate/up and down cases constrain each boundary to 1%. -check.case public @qwen3_moe_expert_mlp_q4k_f16_handoff_differential_case { - %token_count = check.literal value(67) : index - %gate_input_size = check.literal value(512) : index - %gate_output_size = check.literal value(256) : index - %down_output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> - %q8_input = check.generate.fill value(0) : tensor<67x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %partition_table = check.generate.fill value(-1) : tensor<10xi32> - %gate_weight = check.generate.fill value(34) : tensor<4x256x2x144xi8> - %up_weight = check.generate.fill value(35) : tensor<4x256x2x144xi8> - %down_weight = check.generate.fill value(-86) : tensor<4x33x1x144xi8> - %reference_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf32> - %q8_down_input = check.generate.fill value(0) : tensor<67x2x288xi8> - %expected_output = check.generate.fill value(1.0) : tensor<67x33xf32> - %actual_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf16> - %actual_routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual_output = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %gate_input_size](%token_count, %gate_input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %gate_output_size](%token_count, %route_count, %route_stride, %expert_count, %gate_output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %reference_gate_output) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %gate_output_size](%routed_row_count, %gate_output_size, %reference_gate_output, %q8_down_input) : [index, index](index, index, tensor<67x2x256xf32>, tensor<67x2x288xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size](%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size, %q8_down_input, %route_ids, %route_weights, %down_weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x288xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x1x144xi8>, tensor<67x33xf32>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %actual_gate_output) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf16>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %actual_gate_output, %expert_table, %down_weight, %actual_routed_output) : [index](index, tensor<67x2x256xf16>, tensor<272xi32>, tensor<4x33x1x144xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %actual_routed_output, %actual_output) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) - check.expect.close actual(%actual_output) expected(%expected_output) atol(0.5) rtol(0.05) nan(same) : tensor<67x33xf32> +// Compare separate output storage with the in-place residual/output binding +// used by composed routed-down schedules. +check.case public @qwen3_moe_routed_down_weighted_reduce_f16_f32_alias_case { + %token_count = check.literal value(1) : index + %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<1x8xf32> + %routed_output = check.generate.iota offset(-0.25) step(0.015625) period(31) : tensor<1x8x2048xf16> + %separate_residual = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<1x2048xf32> + %in_place_residual = check.generate.iota offset(-0.5) step(0.0078125) period(17) : tensor<1x2048xf32> + %separate_output = check.generate.fill value(0.0) : tensor<1x2048xf32> + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %separate_residual, %separate_output) : [index](index, tensor<1x8xf32>, tensor<1x8x2048xf16>, tensor<1x2048xf32>, tensor<1x2048xf32>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %in_place_residual, %in_place_residual) : [index](index, tensor<1x8xf32>, tensor<1x8x2048xf16>, tensor<1x2048xf32>, tensor<1x2048xf32>) + check.expect.equal actual(%in_place_residual) expected(%separate_output) : tensor<1x2048xf32> check.return } -// Q6_K shares the production matrix schedule but has an independent packed -// decoder and Q8_1 reference provider. -check.case public @qwen3_moe_expert_mlp_q6k_f16_handoff_differential_case { - %token_count = check.literal value(67) : index - %gate_input_size = check.literal value(512) : index - %gate_output_size = check.literal value(256) : index - %down_output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> - %q8_input = check.generate.fill value(0) : tensor<67x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %partition_table = check.generate.fill value(-1) : tensor<10xi32> - %gate_weight = check.generate.fill value(34) : tensor<4x256x2x144xi8> - %up_weight = check.generate.fill value(35) : tensor<4x256x2x144xi8> - %down_weight = check.generate.fill value(-86) : tensor<4x33x1x210xi8> - %reference_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf32> - %q8_down_input = check.generate.fill value(0) : tensor<67x2x288xi8> - %expected_output = check.generate.fill value(1.0) : tensor<67x33xf32> - %actual_gate_output = check.generate.fill value(1.0) : tensor<67x2x256xf16> - %actual_routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual_output = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %gate_input_size](%token_count, %gate_input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %gate_output_size](%token_count, %route_count, %route_stride, %expert_count, %gate_output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %reference_gate_output) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %gate_output_size](%routed_row_count, %gate_output_size, %reference_gate_output, %q8_down_input) : [index, index](index, index, tensor<67x2x256xf32>, tensor<67x2x288xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size](%token_count, %gate_output_size, %route_count, %route_stride, %expert_count, %down_output_size, %q8_down_input, %route_ids, %route_weights, %down_weight, %expected_output) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x288xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x1x210xi8>, tensor<67x33xf32>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %actual_gate_output) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x256x2x144xi8>, tensor<4x256x2x144xi8>, tensor<67x2x256xf16>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %actual_gate_output, %expert_table, %down_weight, %actual_routed_output) : [index](index, tensor<67x2x256xf16>, tensor<272xi32>, tensor<4x33x1x210xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %actual_routed_output, %actual_output) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) - check.expect.close actual(%actual_output) expected(%expected_output) atol(0.5) rtol(0.05) nan(same) : tensor<67x33xf32> - check.return -} - -// Differential coverage crosses token, route, expert, channel, and output-tail -// boundaries. The direct Q8_1 provider supplies an independent raw-Q4_K -// reference while the grouped provider exercises collision-free FP16 -// publication and a separate weighted FP32 residual reduction. -check.case public @qwen3_moe_routed_down_q4k_f16_wmma_differential_case { - %token_count = check.literal value(67) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %input_f32 = check.generate.fill value(0.00390625) : tensor<67x2x512xf32> - %input_f16 = check.generate.fill value(0.00390625) : tensor<67x2x512xf16> - %q8_input = check.generate.fill value(0) : tensor<67x2x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %weight = check.generate.fill value(-86) : tensor<4x33x2x144xi8> - %expected = check.generate.fill value(1.0) : tensor<67x33xf32> - %routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %input_f32, %q8_input) : [index, index](index, index, tensor<67x2x512xf32>, tensor<67x2x576xi8>) - kernel.launch @qwen3_moe_routed_down_q4k_q8_1_x4[%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x576xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x2x144xi8>, tensor<67x33xf32>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input_f16, %expert_table, %weight, %routed_output) : [index](index, tensor<67x2x512xf16>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %actual) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.5) rtol(0.01) nan(same) : tensor<67x33xf32> - check.return -} - -// Q6_K follows the identical routed matrix schedule while selecting its -// independent packed-row decoder. The same irregular shape prevents format -// specialization from weakening routing or tail coverage. -check.case public @qwen3_moe_routed_down_q6k_f16_wmma_differential_case { - %token_count = check.literal value(67) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %routed_row_count = check.literal value(134) : index - %input_f32 = check.generate.fill value(0.00390625) : tensor<67x2x512xf32> - %input_f16 = check.generate.fill value(0.00390625) : tensor<67x2x512xf16> - %q8_input = check.generate.fill value(0) : tensor<67x2x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %route_weights = check.generate.fill value(0.5) : tensor<67x2xf32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %weight = check.generate.fill value(-86) : tensor<4x33x2x210xi8> - %expected = check.generate.fill value(1.0) : tensor<67x33xf32> - %routed_output = check.generate.fill value(0.0) : tensor<67x2x33xf16> - %actual = check.generate.fill value(1.0) : tensor<67x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %input_size](%routed_row_count, %input_size, %input_f32, %q8_input) : [index, index](index, index, tensor<67x2x512xf32>, tensor<67x2x576xi8>) - kernel.launch @qwen3_moe_routed_down_q6k_q8_1_x4[%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size](%token_count, %input_size, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %route_weights, %weight, %expected) : [index, index, index, index, index, index](index, index, index, index, index, index, tensor<67x2x576xi8>, tensor<67x4xi32>, tensor<67x2xf32>, tensor<4x33x2x210xi8>, tensor<67x33xf32>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input_f16, %expert_table, %weight, %routed_output) : [index](index, tensor<67x2x512xf16>, tensor<272xi32>, tensor<4x33x2x210xi8>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %actual) : [index](index, tensor<67x2xf32>, tensor<67x2x33xf16>, tensor<67x33xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.5) rtol(0.01) nan(same) : tensor<67x33xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -// Maximally diverse decode-batch controls. Compact route storage makes the -// flattened 0..127 iota assign every M=16 route to a different expert. -check.case public @qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x144xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q4k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x144xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.case public @qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %input = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - %weight = check.generate.fill value(0) : tensor<128x2048x3x210xi8> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_down_q6k_f16_wmma_grouped[%token_count](%token_count, %input, %expert_table, %weight, %routed_output) : [index](index, tensor<[%token_count]x8x768xf16>, tensor<65664xi32>, tensor<128x2048x3x210xi8>, tensor<[%token_count]x8x2048xf16>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %output) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_differential_case> @qwen3_moe_routed_down_q4k_f16_wmma_small - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_diverse_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q4k_f16_wmma_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_differential_case> @qwen3_moe_routed_down_q6k_f16_wmma_small - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_diverse_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_diverse_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_down_q6k_f16_wmma_benchmark_case> @qwen3_moe_routed_down_q6k_f16_wmma_prefill_512 {token_count = 512} +check.benchmark<@qwen3_moe_routed_down_weighted_reduce_f16_f32_alias_case> @qwen3_moe_routed_down_weighted_reduce_f16_f32_alias 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 f08559f7e55f..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 @@ -5,6 +5,10 @@ // residual results across the workgroup RMS reduction, then publishes both // the updated hidden state and the learned-weight normalized input consumed // 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)] @@ -15,101 +19,27 @@ config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%val config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] -kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer) +kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count$4: index) launch(%token_count$5: index, %route_weights: buffer, %routed_output: buffer, %residual_input: buffer, %output: buffer) + +kernel.decl @qwen3_moe_rmsnorm_f32(%token_count$9: index) launch(%token_count$10: index, %input: buffer, %weight: buffer, %output: buffer) -kernel.decl @qwen3_moe_rmsnorm_f32(%token_count: index) launch(%token_count: 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 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) where [range(%token_count, 1, 2048)] { - %route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index - %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index - %token0 = kernel.workgroup.id : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index + 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.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 - %c8 = index.constant 8 : index - %channel = index.mul %workitem, %c8 : index - %scratch_bytes = index.constant 32 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c0_f32x8 = vector.constant 0.0 : vector<8xf32> - %active_token = index.cmp ult, %token0, %token_count : index - %safe_token0 = scf.select %active_token, %token0, %c0 : index - %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index - %hidden_size_i32 = index.cast %hidden_size : index to i32 - %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 - %assignment_count = index.mul %launch_token_count, %route_count : index - %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer - %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> - %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> - %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %next_projection_input_view = buffer.view %next_projection_input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_view = buffer.view %scratch[%c0_offset] : buffer -> view<8xf32> - - // The predicate depends only on the workgroup ID, so every workitem either - // executes the complete barrier-bearing row reduction or skips it together. - scf.if %active_token { - %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { - %assignment = index.madd %token, %route_count, %route : index - %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> - %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> - %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 - %route_weight_x8 = vector.splat %route_weight : vector<8xf32> - %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> - %next = vector.addf %sum, %weighted : vector<8xf32> - scf.yield %next : vector<8xf32> - } - %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> - %result = vector.addf %residual, %weighted_sum : vector<8xf32> - %squares = vector.mulf %result, %result : vector<8xf32> - %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 - %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 - - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_sum, %scratch_view[%subgroup] : f32, view<8xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index - %is_reduction_lane = index.cmp ult, %lane, %c8 : index - %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 - %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { - %value = view.load %scratch_view[%lane] : view<8xf32> -> f32 - scf.yield %value : f32 - } else { - scf.yield %c0_f32 : f32 - } - %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 - %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 - scf.if %writes_scale { - %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_view[%c0] : f32, view<8xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - - %scale = view.load %scratch_view[%c0] : view<8xf32> -> f32 - %scale_x8 = vector.splat %scale : vector<8xf32> - %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> - %normalized = vector.mulf %result, %scale_x8 : vector<8xf32> - %next_projection = vector.mulf %normalized, %learned_weight : vector<8xf32> - vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> - vector.store %next_projection, %next_projection_input_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> - } + %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 } @@ -124,7 +54,7 @@ check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_differ %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<512x2048xf32> %actual_next_projection_input = check.generate.fill value(1.0) : tensor<512x2048xf32> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state, %expected_hidden_state) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>, tensor<512x2048xf32>) kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %expected_hidden_state, %next_norm_weight, %expected_next_projection_input) : [index](index, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32[%token_count](%token_count, %route_weights, %routed_output, %actual_hidden_state, %next_norm_weight, %actual_next_projection_input) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) check.expect.close actual(%actual_hidden_state) expected(%expected_hidden_state) atol(0.0001) rtol(0.0001) nan(same) : tensor<512x2048xf32> @@ -132,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> @@ -150,7 +127,7 @@ check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_compos %hidden_state = check.generate.fill value(1.0) : tensor<512x2048xf32> %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> %next_projection_input = check.generate.fill value(0.0) : tensor<512x2048xf32> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>) + kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %hidden_state) : [index](index, tensor<512x8xf32>, tensor<512x8x2048xf16>, tensor<512x2048xf32>, tensor<512x2048xf32>) kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<512x2048xf32>, tensor<2048xf32>, tensor<512x2048xf32>) check.return } @@ -158,3 +135,11 @@ check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_compos check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_fused_prefill_512 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/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom deleted file mode 100644 index 00ff052de02c..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom +++ /dev/null @@ -1,269 +0,0 @@ -// Fused grouped-prefill routed residual reduction, next-layer RMSNorm, and -// GGML Q8_1 x4 publication. -// -// One 256-workitem workgroup, comprising eight wave32 subgroups, owns one -// complete 2048-channel token row. Each workitem retains eight adjacent -// residual results across the RMS reduction, updates hidden state, and -// quantizes the learned-weight normalized values directly into the physical -// row consumed by quantized QKV. This is the interlayer producer for the -// quantized attention schedule; no F32 normalized row is materialized. -// -// Four neighboring workitems own one logical 32-element Q8_1 block. Each -// workitem contributes two packed four-value words. LDS preserves the same -// eight-word max and sum reduction order as the standalone GGML packer so the -// fused and decomposed paths produce identical metadata and payload bytes. -amdgpu.target @qwen3_moe_routed_down_next_norm_q8_gfx11_wave32 {subgroup_size = 32} - -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 - -config.decl @qwen3_moe.routed_down.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_down.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @qwen3_moe_routed_down_weighted_reduce_f16_f32(%token_count: index) launch(%token_count: index, %route_weights: buffer, %routed_output: buffer, %hidden_state: buffer) - -kernel.decl @qwen3_moe_rmsnorm_f32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.def target(@qwen3_moe_routed_down_next_norm_q8_gfx11_wave32) @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4(%token_count: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - kernel.launch.config workgroups(%token_capacity, %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) { - %route_count0 = config.get @qwen3_moe.routed_down.route_count : index - %output_size0 = config.get @qwen3_moe.routed_down.output_size : index - %hidden_size0 = config.get @qwen3_moe.model.hidden_size : index - %epsilon = config.get @qwen3_moe.model.rms_epsilon : f32 - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048), le(%token_count, %token_capacity)] : index - %route_count = index.assume %route_count0 [range(%route_count0, 8, 8)] : index - %output_size, %hidden_size = index.assume %output_size0, %hidden_size0 [range(%output_size0, 2048, 2048), range(%hidden_size0, 2048, 2048), eq(%output_size0, %hidden_size0)] : index, index - %token0 = kernel.workgroup.id : index - %workitem0 = kernel.workitem.id : index - %workitem = index.assume %workitem0 [range(%workitem0, 0, 255)] : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c128 = index.constant 128 : index - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %scratch_d_byte_add = index.constant 2048 : offset - %scratch_bytes = index.constant 2304 : offset - %c0_offset = index.constant 0 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127 = scalar.constant 127.0 : f32 - %c0_f32x8 = vector.constant 0.0 : vector<8xf32> - %active_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %active_token, %token0, %c0 : index - %token, %launch_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %hidden_size_i32 = index.cast %hidden_size : index to i32 - %hidden_size_f32 = scalar.sitofp %hidden_size_i32 : i32 to f32 - %assignment_count = index.mul %launch_token_count, %route_count : 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 - %route_weights_noalias, %routed_output_noalias, %hidden_state_noalias, %next_norm_weight_noalias, %next_projection_input_noalias = buffer.assume.noalias %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input : buffer, buffer, buffer, buffer, buffer - %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%route_count]xf32> - %routed_output_view = buffer.view %routed_output_noalias[%c0_offset] : buffer -> view<[%assignment_count]x[%output_size]xf16> - %hidden_state_view = buffer.view %hidden_state_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> - %next_norm_weight_view = buffer.view %next_norm_weight_noalias[%c0_offset] : buffer -> view<[%hidden_size]xf32> - %scratch = buffer.alloca align(16) %scratch_bytes : buffer - %scratch_values = buffer.view %scratch[%c0_offset] : buffer -> view<512xf32> - %scratch_d = buffer.view %scratch[%scratch_d_byte_add] : buffer -> view<64xf32> - // The token predicate is workgroup-uniform, keeping every RMSNorm and Q8_1 - // barrier either active for the whole workgroup or skipped by all workitems. - scf.if %active_token { - %channel = index.mul %workitem, %c8 : index - %weighted_sum = scf.for %route = [%c0 to %route_count step %c1](%sum = %c0_f32x8 : vector<8xf32>) -> (vector<8xf32>) unroll { - %assignment = index.madd %token, %route_count, %route : index - %routed_f16 = vector.load %routed_output_view[%assignment, %channel] : view<[%assignment_count]x[%output_size]xf16> -> vector<8xf16> - %routed = vector.extf %routed_f16 : vector<8xf16> to vector<8xf32> - %route_weight = view.load %route_weights_view[%token, %route] : view<[%launch_token_count]x[%route_count]xf32> -> f32 - %route_weight_x8 = vector.splat %route_weight : vector<8xf32> - %weighted = vector.mulf %routed, %route_weight_x8 : vector<8xf32> - %next = vector.addf %sum, %weighted : vector<8xf32> - scf.yield %next : vector<8xf32> - } - %residual = vector.load %hidden_state_view[%token, %channel] : view<[%launch_token_count]x[%hidden_size]xf32> -> vector<8xf32> - %result = vector.addf %residual, %weighted_sum : vector<8xf32> - %squares = vector.mulf %result, %result : vector<8xf32> - %thread_sum = vector.reduce %squares, %c0_f32 : vector<8xf32>, f32 - %subgroup_sum = kernel.subgroup.reduce %thread_sum : f32 - - %is_subgroup_leader = index.cmp eq, %lane, %c0 : index - scf.if %is_subgroup_leader { - view.store %subgroup_sum, %scratch_values[%subgroup] : f32, view<512xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %is_reduction_subgroup = index.cmp eq, %subgroup, %c0 : index - %is_reduction_lane = index.cmp ult, %lane, %c8 : index - %loads_subgroup_sum = scalar.andi %is_reduction_subgroup, %is_reduction_lane : i1 - %subgroup_partial = scf.if %loads_subgroup_sum -> (f32) { - %value = view.load %scratch_values[%lane] : view<512xf32> -> f32 - scf.yield %value : f32 - } else { - scf.yield %c0_f32 : f32 - } - %row_sum = kernel.subgroup.reduce %subgroup_partial : f32 - %writes_scale = scalar.andi %is_reduction_subgroup, %is_subgroup_leader : i1 - scf.if %writes_scale { - %mean = scalar.divf %row_sum, %hidden_size_f32 : f32 - %biased_mean = scalar.addf %mean, %epsilon : f32 - %scale = scalar.rsqrtf %biased_mean : f32 - view.store %scale, %scratch_values[%c0] : f32, view<512xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - - %row_scale = view.load %scratch_values[%c0] : view<512xf32> -> f32 - // All waves must capture the row scale before scratch is reused for the - // independent Q8 block reductions. - kernel.barrier scope(workgroup) ordering(acq_rel) - %row_scale_vector = vector.splat %row_scale : vector<8xf32> - %learned_weight = vector.load %next_norm_weight_view[%channel] : view<[%hidden_size]xf32> -> vector<8xf32> - %normalized0 = vector.mulf %result, %row_scale_vector : vector<8xf32> - %normalized = vector.mulf %normalized0, %learned_weight : vector<8xf32> - vector.store %result, %hidden_state_view[%token, %channel] : vector<8xf32>, view<[%launch_token_count]x[%hidden_size]xf32> - - %low_values = vector.slice %normalized[0] : vector<8xf32> -> vector<4xf32> - %high_values = vector.slice %normalized[4] : vector<8xf32> -> vector<4xf32> - %low_absolute_values = vector.absf %low_values : vector<4xf32> - %high_absolute_values = vector.absf %high_values : vector<4xf32> - %low_max = vector.reduce %low_absolute_values, %c0_f32 : vector<4xf32>, f32 - %high_max = vector.reduce %high_absolute_values, %c0_f32 : vector<4xf32>, f32 - %partial_base = index.mul %workitem, %c2 : index - %partial_high = index.add %partial_base, %c1 : index - view.store %low_max, %scratch_values[%partial_base] : f32, view<512xf32> - view.store %high_max, %scratch_values[%partial_high] : f32, view<512xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - - %block_in_row = index.div %workitem, %c4 : index - %workitem_in_block = index.rem %workitem, %c4 : index - %is_block_leader = index.cmp eq, %workitem_in_block, %c0 : index - scf.if %is_block_leader { - %block_partial_base = index.mul %block_in_row, %c8 : index - %block_maxima = vector.load %scratch_values[%block_partial_base] : view<512xf32> -> vector<8xf32> - %amax = vector.reduce %block_maxima, %c0_f32 : vector<8xf32>, f32 - %d = scalar.divf %amax, %c127 : f32 - view.store %d, %scratch_d[%block_in_row] : f32, view<64xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - - %d = view.load %scratch_d[%block_in_row] : view<64xf32> -> f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<8xf32> - %scaled_values = vector.mulf %normalized, %d_inverse_vector : vector<8xf32> - %rounded_values = vector.roundf %scaled_values : vector<8xf32> - %quantized_values = vector.fptosi %rounded_values : vector<8xf32> to vector<8xi8> - %packed_words = vector.bitcast %quantized_values : vector<8xi8> to vector<2xi32> - %physical_group = index.div %block_in_row, %c4 : index - %block_in_group = index.rem %block_in_row, %c4 : index - %group_byte_add = index.scale %physical_group, %group_bytes : index, offset -> offset - %group_byte_offset = index.add %token_output_byte_base, %group_byte_add : offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %next_projection_input_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %next_projection_input_noalias[%payload_byte_offset] : buffer -> view<32xi32> - %block_word_base = index.mul %block_in_group, %c8 : index - %workitem_word_add = index.mul %workitem_in_block, %c2 : index - %packed_word_index0 = index.add %block_word_base, %workitem_word_add : index - %packed_word_index = index.assume %packed_word_index0 [range(%packed_word_index0, 0, 31)] : index - vector.store %packed_words, %group_qs[%packed_word_index] : vector<2xi32>, view<32xi32> - - %low_rounded_values = vector.slice %rounded_values[0] : vector<8xf32> -> vector<4xf32> - %high_rounded_values = vector.slice %rounded_values[4] : vector<8xf32> -> vector<4xf32> - %low_quantized_sum = vector.reduce %low_rounded_values, %c0_f32 : vector<4xf32>, f32 - %high_quantized_sum = vector.reduce %high_rounded_values, %c0_f32 : vector<4xf32>, f32 - view.store %low_quantized_sum, %scratch_values[%partial_base] : f32, view<512xf32> - view.store %high_quantized_sum, %scratch_values[%partial_high] : f32, view<512xf32> - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.if %is_block_leader { - %block_partial_base = index.mul %block_in_row, %c8 : index - %block_sums = vector.load %scratch_values[%block_partial_base] : view<512xf32> -> vector<8xf32> - %quantized_sum = vector.reduce %block_sums, %c0_f32 : vector<8xf32>, f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block_in_group, %c2 : index - %s_index = index.add %ds_index, %c1 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - } - kernel.return -} - -// Fourteen independent token rows lock the complete interlayer contract. The -// fused producer must match ordinary weighted reduction for hidden state and -// RMSNorm followed by the generic GGML packer for every Q8 metadata and -// payload byte. -check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_differential_case { - %token_count = check.literal value(14) : index - %hidden_size = check.literal value(2048) : index - %route_weights = check.generate.iota offset(0.03125) step(0.015625) period(8) : tensor<14x8xf32> - %routed_seed = check.param.seed base(5858425849430430008) count(1) : i64 - %routed_output = check.generate.random.uniform seed(%routed_seed) range(-0.5 to 0.5) : tensor<14x8x2048xf16> - %expected_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(31) : tensor<14x2048xf32> - %actual_hidden_state = check.generate.iota offset(-0.5) step(0.0078125) period(31) : tensor<14x2048xf32> - %next_norm_weight = check.generate.iota offset(0.5) step(0.0078125) period(29) : tensor<2048xf32> - %normalized = check.generate.fill value(0.0) : tensor<14x2048xf32> - %expected_q8 = check.generate.fill value(0) : tensor<14x2304xi8> - %actual_q8 = check.generate.fill value(1) : tensor<14x2304xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %expected_hidden_state) : [index](index, tensor<14x8xf32>, tensor<14x8x2048xf16>, tensor<14x2048xf32>) - kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %expected_hidden_state, %next_norm_weight, %normalized) : [index](index, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2048xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %expected_q8) : [index, index](index, index, tensor<14x2048xf32>, tensor<14x2304xi8>) - kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4[%token_count](%token_count, %route_weights, %routed_output, %actual_hidden_state, %next_norm_weight, %actual_q8) : [index](index, tensor<14x8xf32>, tensor<14x8x2048xf16>, tensor<14x2048xf32>, tensor<2048xf32>, tensor<14x2304xi8>) - check.expect.close actual(%actual_hidden_state) expected(%expected_hidden_state) atol(0.0001) rtol(0.0001) nan(same) : tensor<14x2048xf32> - check.expect.equal actual(%actual_q8) expected(%expected_q8) : tensor<14x2304xi8> - check.return -} - -check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case { - %token_count = check.param.choice values([14, 32, 128, 512]) name("token_count") : index - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %hidden_state = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %next_projection_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4[%token_count](%token_count, %route_weights, %routed_output, %hidden_state, %next_norm_weight, %next_projection_input) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2304xi8>) - check.return -} - -check.case public @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case { - %token_count = check.param.choice values([14, 32, 128, 512]) name("token_count") : index - %hidden_size = check.literal value(2048) : index - %route_weights = check.generate.fill value(0.125) : tensor<[%token_count]x8xf32> - %routed_output = check.generate.fill value(0.0) : tensor<[%token_count]x8x2048xf16> - %hidden_state = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %next_norm_weight = check.generate.fill value(1.0) : tensor<2048xf32> - %normalized = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %next_projection_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - kernel.launch @qwen3_moe_routed_down_weighted_reduce_f16_f32[%token_count](%token_count, %route_weights, %routed_output, %hidden_state) : [index](index, tensor<[%token_count]x8xf32>, tensor<[%token_count]x8x2048xf16>, tensor<[%token_count]x2048xf32>) - kernel.launch @qwen3_moe_rmsnorm_f32[%token_count](%token_count, %hidden_state, %next_norm_weight, %normalized) : [index](index, tensor<[%token_count]x2048xf32>, tensor<2048xf32>, tensor<[%token_count]x2048xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %hidden_size](%token_count, %hidden_size, %normalized, %next_projection_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) - check.return -} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_14 {token_count = 14} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_fused_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_prefill_14 {token_count = 14} - -check.benchmark<@qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_benchmark_case> @qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_q8_1_x4_composed_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom deleted file mode 100644 index ce82133b44aa..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_gate_up_swiglu_q4k.loom +++ /dev/null @@ -1,1661 +0,0 @@ -// Fuses the Qwen3 MoE gate and up expert projections with the SwiGLU -// epilogue. Weights are consumed directly from GGUF's raw Q4_K layout: -// [expert][output channel][K / 256][144 bytes]. Activations use GGML's Q8_1 -// x4 layout produced by the shared quantization kernel. -// -// Route IDs have a logical route_count and an independent physical -// route_stride. This admits llama.cpp's [token][128] argsort storage while -// operating on only the selected top-8 entries, and also admits compact route -// buffers produced by a future fused router. -// -// Two schedules share this contract. Small token counts use one wave per -// [token, route, output channel]. Larger token counts first build a compact -// assignment table per expert, then group up to 32 selected rows with 32 -// output channels so raw weights and quantized activations can be reused -// through workgroup memory. -config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] - -config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] - -config.decl @qwen3_moe.workload.token_capacity : %value: index where [range(%value, 1, 2048)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -func.decl @ggml_q8_1_x4_block(%q8_input: buffer, %row_byte_base: offset, %q8_block: index) -> (vector<32xi8>, f32, f32) - -// Decodes one 16-element Q4_K chunk while preserving its unsigned nibbles for -// AMDGPU's mixed u8*s8 dot4 instruction. The returned scale terms are shared -// by every routed activation row consuming the same weight chunk. -func.def inline @qwen3_moe_q4k_chunk_local(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index) -> (vector<16xi8>, f32, f32) { - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 144 : offset - %header_byte_count = index.constant 4 : index - %code_byte_add = index.constant 16 : offset - %c4_i32 = scalar.constant 4 : i32 - %c6_i32 = scalar.constant 6 : i32 - %c15_i32 = scalar.constant 15 : i32 - %c63_i32 = scalar.constant 63 : i32 - %nibble_mask = vector.constant 252645135 : vector<4xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_byte_add : offset - %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %header_words = vector.load %header_view[0] : view<4xi32> -> vector<4xi32> - %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> - %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 - %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %header_bytes = vector.bitcast %header_words : vector<4xi32> to vector<16xi8> - %iqs_group = index.mul %bounded_group, %c8 : index - %iqs_half = index.mul %bounded_half, %c4 : index - %iqs = index.add %iqs_group, %iqs_half : index - %qs_page0 = index.div %iqs, %c16 : index - %qs_page = index.mul %qs_page0, %c8 : index - %qs_lane = index.rem %iqs, %c8 : index - %qs_index0 = index.add %qs_page, %qs_lane : index - %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 28)] : index - %iqs_mod16 = index.rem %iqs, %c16 : index - %nibble_page = index.div %iqs_mod16, %c8 : index - %nibble_shift_index = index.mul %nibble_page, %c4 : index - %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 - %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> - %packed_codes = vector.load %code_view[%qs_index] : view<32xi32> -> vector<4xi32> - %shifted_codes = vector.shrui %packed_codes, %nibble_shift : vector<4xi32> - %masked_codes = vector.andi %shifted_codes, %nibble_mask : vector<4xi32> - %q4_values = vector.bitcast %masked_codes : vector<4xi32> to vector<16xi8> - %is_low_group = index.cmp ult, %bounded_group, %c4 : index - %scale, %minimum = scf.if %is_low_group -> (i32, i32) { - %scale_index = index.add %bounded_group, %header_byte_count : index - %minimum_index0 = index.add %bounded_group, %c4 : index - %minimum_index = index.add %minimum_index0, %header_byte_count : index - %scale_i8 = vector.extract %header_bytes[%scale_index] : vector<16xi8> -> i8 - %minimum_i8 = vector.extract %header_bytes[%minimum_index] : vector<16xi8> -> i8 - %scale_u8 = scalar.extui %scale_i8 : i8 to i32 - %minimum_u8 = scalar.extui %minimum_i8 : i8 to i32 - %scale_low6 = scalar.andi %scale_u8, %c63_i32 : i32 - %minimum_low6 = scalar.andi %minimum_u8, %c63_i32 : i32 - scf.yield %scale_low6, %minimum_low6 : i32, i32 - } else { - %packed_index0 = index.add %bounded_group, %c4 : index - %packed_index = index.add %packed_index0, %header_byte_count : index - %scale_high_index = index.sub %bounded_group, %c4 : index - %scale_high_header_index = index.add %scale_high_index, %header_byte_count : index - %minimum_high_header_index = index.add %bounded_group, %header_byte_count : index - %packed_i8 = vector.extract %header_bytes[%packed_index] : vector<16xi8> -> i8 - %scale_high_i8 = vector.extract %header_bytes[%scale_high_header_index] : vector<16xi8> -> i8 - %minimum_high_i8 = vector.extract %header_bytes[%minimum_high_header_index] : vector<16xi8> -> i8 - %packed = scalar.extui %packed_i8 : i8 to i32 - %scale_high = scalar.extui %scale_high_i8 : i8 to i32 - %minimum_high = scalar.extui %minimum_high_i8 : i8 to i32 - %scale_low4 = scalar.andi %packed, %c15_i32 : i32 - %minimum_low4 = scalar.shrui %packed, %c4_i32 : i32 - %scale_high2_raw = scalar.shrui %scale_high, %c6_i32 : i32 - %minimum_high2_raw = scalar.shrui %minimum_high, %c6_i32 : i32 - %scale_high2 = scalar.shli %scale_high2_raw, %c4_i32 : i32 - %minimum_high2 = scalar.shli %minimum_high2_raw, %c4_i32 : i32 - %scale_value = scalar.ori %scale_low4, %scale_high2 : i32 - %minimum_value = scalar.ori %minimum_low4, %minimum_high2 : i32 - scf.yield %scale_value, %minimum_value : i32, i32 - } - %scale_f32 = scalar.uitofp %scale : i32 to f32 - %minimum_f32 = scalar.uitofp %minimum : i32 to f32 - %d_scale = scalar.mulf %d, %scale_f32 : f32 - %dmin_scale = scalar.mulf %dmin, %minimum_f32 : f32 - func.return %q4_values, %d_scale, %dmin_scale : vector<16xi8>, f32, f32 -} - -// Global-weight variant of the Q4_K chunk decoder. Scalar scale-byte loads are -// legal and cheaper for the one-wave schedule, while the LDS schedule uses the -// packed-header variant above because AMDGPU has no local i8 load descriptor. -func.def inline @qwen3_moe_q4k_chunk_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index) -> (vector<16xi8>, f32, f32) { - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %block_bytes = index.constant 144 : offset - %scale_byte_add = index.constant 4 : offset - %code_byte_add = index.constant 16 : offset - %c4_i32 = scalar.constant 4 : i32 - %c6_i32 = scalar.constant 6 : i32 - %c15_i32 = scalar.constant 15 : i32 - %c63_i32 = scalar.constant 63 : i32 - %nibble_mask = vector.constant 252645135 : vector<4xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_byte_add : offset - %half_view = buffer.view %weight[%block_byte_base] : buffer -> view<2xf16> - %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<12xi8> - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %d_f16 = view.load %half_view[0] : view<2xf16> -> f16 - %dmin_f16 = view.load %half_view[1] : view<2xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %iqs_group = index.mul %bounded_group, %c8 : index - %iqs_half = index.mul %bounded_half, %c4 : index - %iqs = index.add %iqs_group, %iqs_half : index - %qs_page0 = index.div %iqs, %c16 : index - %qs_page = index.mul %qs_page0, %c8 : index - %qs_lane = index.rem %iqs, %c8 : index - %qs_index0 = index.add %qs_page, %qs_lane : index - %qs_index = index.assume %qs_index0 [range(%qs_index0, 0, 28)] : index - %iqs_mod16 = index.rem %iqs, %c16 : index - %nibble_page = index.div %iqs_mod16, %c8 : index - %nibble_shift_index = index.mul %nibble_page, %c4 : index - %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 - %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> - %packed_codes = vector.load %code_view[%qs_index] : view<32xi32> -> vector<4xi32> - %shifted_codes = vector.shrui %packed_codes, %nibble_shift : vector<4xi32> - %masked_codes = vector.andi %shifted_codes, %nibble_mask : vector<4xi32> - %q4_values = vector.bitcast %masked_codes : vector<4xi32> to vector<16xi8> - %is_low_group = index.cmp ult, %bounded_group, %c4 : index - %scale, %minimum = scf.if %is_low_group -> (i32, i32) { - %minimum_index = index.add %bounded_group, %c4 : index - %scale_i8 = view.load %scale_view[%bounded_group] : view<12xi8> -> i8 - %minimum_i8 = view.load %scale_view[%minimum_index] : view<12xi8> -> i8 - %scale_u8 = scalar.extui %scale_i8 : i8 to i32 - %minimum_u8 = scalar.extui %minimum_i8 : i8 to i32 - %scale_low6 = scalar.andi %scale_u8, %c63_i32 : i32 - %minimum_low6 = scalar.andi %minimum_u8, %c63_i32 : i32 - scf.yield %scale_low6, %minimum_low6 : i32, i32 - } else { - %packed_index = index.add %bounded_group, %c4 : index - %scale_high_index = index.sub %bounded_group, %c4 : index - %packed_i8 = view.load %scale_view[%packed_index] : view<12xi8> -> i8 - %scale_high_i8 = view.load %scale_view[%scale_high_index] : view<12xi8> -> i8 - %minimum_high_i8 = view.load %scale_view[%bounded_group] : view<12xi8> -> i8 - %packed = scalar.extui %packed_i8 : i8 to i32 - %scale_high = scalar.extui %scale_high_i8 : i8 to i32 - %minimum_high = scalar.extui %minimum_high_i8 : i8 to i32 - %scale_low4 = scalar.andi %packed, %c15_i32 : i32 - %minimum_low4 = scalar.shrui %packed, %c4_i32 : i32 - %scale_high2_raw = scalar.shrui %scale_high, %c6_i32 : i32 - %minimum_high2_raw = scalar.shrui %minimum_high, %c6_i32 : i32 - %scale_high2 = scalar.shli %scale_high2_raw, %c4_i32 : i32 - %minimum_high2 = scalar.shli %minimum_high2_raw, %c4_i32 : i32 - %scale_value = scalar.ori %scale_low4, %scale_high2 : i32 - %minimum_value = scalar.ori %minimum_low4, %minimum_high2 : i32 - scf.yield %scale_value, %minimum_value : i32, i32 - } - %scale_f32 = scalar.uitofp %scale : i32 to f32 - %minimum_f32 = scalar.uitofp %minimum : i32 to f32 - %d_scale = scalar.mulf %d, %scale_f32 : f32 - %dmin_scale = scalar.mulf %dmin, %minimum_f32 : f32 - func.return %q4_values, %d_scale, %dmin_scale : vector<16xi8>, f32, f32 -} - -// Decodes one Q4_K scale/minimum pair from a header loaded as three packed -// scale words. Keeping the complete header in registers lets paired-group -// contractions share both the header and packed-code load. -func.def inline @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) { - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c2_i32 = scalar.constant 2 : i32 - %c4_i32 = scalar.constant 4 : i32 - %c15_i32 = scalar.constant 15 : i32 - %c48_i32 = scalar.constant 48 : i32 - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %is_low_group = index.cmp ult, %bounded_group, %c4 : index - %scale_lane = index.rem %bounded_group, %c4 : index - %scale_shift_index = index.mul %scale_lane, %c8 : index - %scale_shift = index.cast %scale_shift_index : index to i32 - %high_shift = scalar.addi %scale_shift, %c2_i32 : i32 - %minimum_shift = scalar.addi %scale_shift, %c4_i32 : i32 - %selected_scale_source = scf.select %is_low_group, %scale0, %scale2 : i32 - %selected_minimum_source = scf.select %is_low_group, %scale1, %scale2 : i32 - %selected_scale_high_shift = scf.select %is_low_group, %scale_shift, %high_shift : i32 - %selected_minimum_low_shift = scf.select %is_low_group, %scale_shift, %minimum_shift : i32 - %scale_low0 = scalar.shrui %selected_scale_source, %scale_shift : i32 - %scale_low = scalar.andi %scale_low0, %c15_i32 : i32 - %scale_high0 = scalar.shrui %scale0, %selected_scale_high_shift : i32 - %scale_high = scalar.andi %scale_high0, %c48_i32 : i32 - %scale = scalar.ori %scale_low, %scale_high : i32 - %minimum_low0 = scalar.shrui %selected_minimum_source, %selected_minimum_low_shift : i32 - %minimum_low = scalar.andi %minimum_low0, %c15_i32 : i32 - %minimum_high0 = scalar.shrui %scale1, %selected_scale_high_shift : i32 - %minimum_high = scalar.andi %minimum_high0, %c48_i32 : i32 - %minimum = scalar.ori %minimum_low, %minimum_high : i32 - func.return %scale, %minimum : i32, i32 -} - -// Decodes two adjacent Q4_K groups from the low and high nibbles of one -// packed-code load. The caller supplies the already-loaded 16-byte header. -func.def inline @qwen3_moe_q4k_chunk_pair_global(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %q4_half: index, %header_words: vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) { - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %block_bytes = index.constant 144 : offset - %code_byte_add = index.constant 16 : offset - %c4_i32 = scalar.constant 4 : i32 - %nibble_mask = vector.constant 252645135 : vector<4xi32> - %bounded_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index - %bounded_half = index.assume %q4_half [range(%q4_half, 0, 1)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_byte_add : offset - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> - %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 - %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 - %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 - %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 - %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %pair_code_base = index.mul %bounded_pair, %c8 : index - %half_code_add = index.mul %bounded_half, %c4 : index - %code_index0 = index.add %pair_code_base, %half_code_add : index - %code_index = index.assume %code_index0 [range(%code_index0, 0, 28)] : index - %packed_codes = vector.load %code_view[%code_index] : view<32xi32> -> vector<4xi32> - %low_codes = vector.andi %packed_codes, %nibble_mask : vector<4xi32> - %c4_i32v = vector.splat %c4_i32 : vector<4xi32> - %high_shifted = vector.shrui %packed_codes, %c4_i32v : vector<4xi32> - %high_codes = vector.andi %high_shifted, %nibble_mask : vector<4xi32> - %q4_low = vector.bitcast %low_codes : vector<4xi32> to vector<16xi8> - %q4_high = vector.bitcast %high_codes : vector<4xi32> to vector<16xi8> - %low_group = index.mul %bounded_pair, %c2 : index - %high_group = index.add %low_group, %c1 : index - %low_scale, %low_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %low_group) : (i32, i32, i32, index) -> (i32, i32) - %high_scale, %high_minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %high_group) : (i32, i32, i32, index) -> (i32, i32) - %low_scale_f32 = scalar.uitofp %low_scale : i32 to f32 - %low_minimum_f32 = scalar.uitofp %low_minimum : i32 to f32 - %high_scale_f32 = scalar.uitofp %high_scale : i32 to f32 - %high_minimum_f32 = scalar.uitofp %high_minimum : i32 to f32 - %low_d_scale = scalar.mulf %d, %low_scale_f32 : f32 - %low_dmin_scale = scalar.mulf %dmin, %low_minimum_f32 : f32 - %high_d_scale = scalar.mulf %d, %high_scale_f32 : f32 - %high_dmin_scale = scalar.mulf %dmin, %high_minimum_f32 : f32 - func.return %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale : vector<16xi8>, f32, f32, vector<16xi8>, f32, f32 -} - -// Contracts one decoded Q4_K half-group with a Q8_1 half-block. Both -// half-groups apply half of the Q8_1 block-sum correction. -func.def inline @qwen3_moe_q4k_q8_1_dot(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { - %c0_i32 = scalar.constant 0 : i32 - %c0_i32v = vector.constant 0 : vector<4xi32> - %half_f32 = scalar.constant 0.5 : f32 - %partial_dots = vector.dot4i %q4_values, %q8_values, %c0_i32v : vector<16xi8>, vector<16xi8>, vector<4xi32> - %q_sum = vector.reduce %partial_dots, %c0_i32 : vector<4xi32>, i32 - %q_sum_f32 = scalar.sitofp %q_sum : i32 to f32 - %scaled_dot0 = scalar.mulf %q8_d, %d_scale : f32 - %scaled_dot = scalar.mulf %scaled_dot0, %q_sum_f32 : f32 - %q8_half_sum = scalar.mulf %q8_s, %half_f32 : f32 - %minimum_correction = scalar.mulf %dmin_scale, %q8_half_sum : f32 - %contribution = scalar.subf %scaled_dot, %minimum_correction : f32 - func.return %contribution : f32 -} - -// Contracts one decoded Q4_K half-group against four routed Q8_1 rows. The -// leading row axis stays explicit through dot4 and scale correction so the -// grouped schedule can reuse each weight decode without scalarizing the rows. -func.def inline @qwen3_moe_q4k_q8_1_dot4_rows(%q4_values: vector<16xi8>, %d_scale: f32, %dmin_scale: f32, %q8_values: vector<4x16xi8>, %q8_d: vector<4xf32>, %q8_s: vector<4xf32>) -> (vector<4xf32>) { - %q8_values0 = vector.extract %q8_values[0] : vector<4x16xi8> -> vector<16xi8> - %q8_values1 = vector.extract %q8_values[1] : vector<4x16xi8> -> vector<16xi8> - %q8_values2 = vector.extract %q8_values[2] : vector<4x16xi8> -> vector<16xi8> - %q8_values3 = vector.extract %q8_values[3] : vector<4x16xi8> -> vector<16xi8> - %q8_d0 = vector.extract %q8_d[0] : vector<4xf32> -> f32 - %q8_d1 = vector.extract %q8_d[1] : vector<4xf32> -> f32 - %q8_d2 = vector.extract %q8_d[2] : vector<4xf32> -> f32 - %q8_d3 = vector.extract %q8_d[3] : vector<4xf32> -> f32 - %q8_s0 = vector.extract %q8_s[0] : vector<4xf32> -> f32 - %q8_s1 = vector.extract %q8_s[1] : vector<4xf32> -> f32 - %q8_s2 = vector.extract %q8_s[2] : vector<4xf32> -> f32 - %q8_s3 = vector.extract %q8_s[3] : vector<4xf32> -> f32 - %contribution0 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values0, %q8_d0, %q8_s0) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution1 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values1, %q8_d1, %q8_s1) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution2 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values2, %q8_d2, %q8_s2) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution3 = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values3, %q8_d3, %q8_s3) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %contribution = vector.from_elements %contribution0, %contribution1, %contribution2, %contribution3 : vector<4xf32> - func.return %contribution : vector<4xf32> -} - -// Convenience wrapper used by the one-wave schedule. -func.def inline @qwen3_moe_q4k_q8_1_chunk(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %q4_half: index, %q8_values: vector<16xi8>, %q8_d: f32, %q8_s: f32) -> (f32) { - %q4_values, %d_scale, %dmin_scale = func.call @qwen3_moe_q4k_chunk_global(%weight, %row_byte_base, %q4_block, %q4_group, %q4_half) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %contribution = func.call @qwen3_moe_q4k_q8_1_dot(%q4_values, %d_scale, %dmin_scale, %q8_values, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - func.return %contribution : f32 -} - -// Computes one lane's partial for a complete Q4_K row. Dense and routed -// providers choose output ownership independently, then share this exact -// packed-row contraction before reducing across a wave32 subgroup. -func.def inline @qwen3_moe_q4k_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, 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 - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c128 = index.constant 128 : index - %c144 = index.constant 144 : index - %c256 = index.constant 256 : index - %c511 = index.constant 511 : index - %c512 = index.constant 512 : index - %q8_group_bytes = index.constant 144 : offset - %q8_payload_byte_add = index.constant 16 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %q8_group_count = index.div %bounded_input_size, %c128 : index - %padded_input_size = index.add %bounded_input_size, %c511 : index - %iteration_count = index.div %padded_input_size, %c512 : index - %q4_group0 = index.div %bounded_lane, %c2 : index - %q4_group1 = index.rem %q4_group0, %c8 : index - %q4_group = index.assume %q4_group1 [range(%q4_group1, 0, 7)] : index - %q4_half0 = index.rem %bounded_lane, %c2 : index - %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index - %lane_q4_block = index.div %bounded_lane, %c16 : index - %lane_q8_group = index.div %bounded_lane, %c8 : index - %q8_inner_block0 = index.div %bounded_lane, %c2 : index - %q8_inner_block1 = index.rem %q8_inner_block0, %c4 : index - %q8_inner_block = index.assume %q8_inner_block1 [range(%q8_inner_block1, 0, 3)] : index - %q8_inner_word_base = index.mul %q8_inner_block, %c8 : index - %q8_half_word_add = index.mul %q4_half, %c4 : index - %q8_word_index0 = index.add %q8_inner_word_base, %q8_half_word_add : index - %q8_word_index = index.assume %q8_word_index0 [range(%q8_word_index0, 0, 28)] : index - %q8_ds_index = index.mul %q8_inner_block, %c2 : index - %q8_s_index = index.add %q8_ds_index, %c1 : index - %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { - %iteration_q4_block = index.mul %iteration, %c2 : index - %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index - %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index - %contribution = scf.if %valid_q4_block -> (f32) { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %iteration_q8_group = index.mul %iteration, %c4 : index - %q8_group0 = index.add %iteration_q8_group, %lane_q8_group : index - %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index - %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset - %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset - %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset - %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> - %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> - %q8_d_f16 = view.load %q8_ds_view[%q8_ds_index] : view<8xf16> -> f16 - %q8_s_f16 = view.load %q8_ds_view[%q8_s_index] : view<8xf16> -> f16 - %q8_d = scalar.extf %q8_d_f16 : f16 to f32 - %q8_s = scalar.extf %q8_s_f16 : f16 to f32 - %q8_words = vector.load %q8_words_view[%q8_word_index] : view<32xi32> -> vector<4xi32> - %q8_values = vector.bitcast %q8_words : vector<4xi32> to vector<16xi8> - %dot = func.call @qwen3_moe_q4k_q8_1_chunk(%weight, %weight_row_byte_base, %q4_block, %q4_group, %q4_half, %q8_values, %q8_d, %q8_s) : (buffer, offset, index, index, index, vector<16xi8>, f32, f32) -> (f32) - scf.yield %dot : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %iteration_acc, %contribution : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Computes one lane's two adjacent-group contributions within one Q4_K block. -// Eight block lanes consume both nibbles from every packed code load and cover -// the block's full 256-element input extent. -func.def inline @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %q4_block: index, %block_lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_q4_block0 = index.assume %q4_block [range(%q4_block, 0, 127)] : index - %bounded_block_lane = index.assume %block_lane [range(%block_lane, 0, 7)] : 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 - %c256 = index.constant 256 : index - %q4_block_bytes = index.constant 144 : offset - %q8_group_bytes = index.constant 144 : offset - %q8_payload_byte_add = index.constant 16 : offset - %q4_block_count = index.div %bounded_input_size, %c256 : index - %q8_group_count = index.div %bounded_input_size, %c128 : index - %bounded_q4_block, %bounded_q4_block_count = index.assume %bounded_q4_block0, %q4_block_count [lt(%bounded_q4_block0, %q4_block_count)] : index, index - %q4_group_pair0 = index.div %bounded_block_lane, %c2 : index - %q4_group_pair = index.assume %q4_group_pair0 [range(%q4_group_pair0, 0, 3)] : index - %q4_half0 = index.rem %bounded_block_lane, %c2 : index - %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index - %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index - %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index - %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index - %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index - %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index - %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index - %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index - %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index - %q8_half_word_add = index.mul %q4_half, %c4 : index - %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index - %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index - %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index - %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index - %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index - %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index - %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index - %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index - %q8_block_group_base = index.mul %bounded_q4_block, %c2 : index - %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index - %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index - %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset - %q8_group_byte_base = index.add %q8_row_byte_base, %q8_group_byte_add : offset - %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset - %q8_ds_view = buffer.view %q8_input[%q8_group_byte_base] : buffer -> view<8xf16> - %q8_words_view = buffer.view %q8_input[%q8_payload_byte_base] : buffer -> view<32xi32> - %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> - %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 - %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 - %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 - %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 - %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 - %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 - %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 - %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 - %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> - %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> - %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> - %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> - %q4_block_byte_add = index.scale %bounded_q4_block, %q4_block_bytes : index, offset -> offset - %q4_block_byte_base = index.add %weight_row_byte_base, %q4_block_byte_add : offset - %q4_header_view = buffer.view %weight[%q4_block_byte_base] : buffer -> view<4xi32> - %q4_header_words = vector.load %q4_header_view[0] : view<4xi32> -> vector<4xi32> - %q4_low, %low_d_scale, %low_dmin_scale, %q4_high, %high_d_scale, %high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%weight, %weight_row_byte_base, %bounded_q4_block, %q4_group_pair, %q4_half, %q4_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) - %low = func.call @qwen3_moe_q4k_q8_1_dot(%q4_low, %low_d_scale, %low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %high = func.call @qwen3_moe_q4k_q8_1_dot(%q4_high, %high_d_scale, %high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %pair = scalar.addf %low, %high : f32 - func.return %pair : f32 -} - -// Computes one lane's partial while consuming both nibbles of each packed -// Q4_K code load. Eight lanes cover one 256-element Q4_K block, so a wave32 -// advances through four blocks per iteration and reduces the adjacent-group -// contributions together. -func.def inline @qwen3_moe_q4k_q8_1_x4_paired_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, 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 - %c1023 = index.constant 1023 : index - %c1024 = index.constant 1024 : index - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %padded_input_size = index.add %bounded_input_size, %c1023 : index - %iteration_count = index.div %padded_input_size, %c1024 : index - %lane_q4_block = index.div %bounded_lane, %c8 : index - %block_lane0 = index.rem %bounded_lane, %c8 : index - %block_lane = index.assume %block_lane0 [range(%block_lane0, 0, 7)] : index - %sum = scf.for %iteration = [%c0 to %iteration_count step %c1](%iteration_acc = %c0_f32 : f32) -> (f32) unroll { - %iteration_q4_block = index.mul %iteration, %c4 : index - %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index - %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index - %contribution = scf.if %valid_q4_block -> (f32) { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %block_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) - scf.yield %pair : f32 - } else { - scf.yield %c0_f32 : f32 - } - %next = scalar.addf %iteration_acc, %contribution : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Computes one lane's partial for a row owned by an eight-lane cohort. Every -// cohort lane consumes both nibbles for one adjacent Q4_K group pair while the -// cohort walks all blocks in the row. This schedule lets separate cohorts -// contract independent routed rows concurrently. -func.def inline @qwen3_moe_q4k_q8_1_x4_cohort_row_lane(%input_size: index, %weight: buffer, %weight_row_byte_base: offset, %q8_input: buffer, %q8_row_byte_base: offset, %cohort_lane: index) -> (f32) { - %bounded_input_size = index.assume %input_size [range(%input_size, 256, 32768), mul(%input_size, 256)] : index - %bounded_cohort_lane = index.assume %cohort_lane [range(%cohort_lane, 0, 7)] : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %q4_block_count = index.div %bounded_input_size, %c256 : index - %sum = scf.for %q4_block0 = [%c0 to %q4_block_count step %c1](%block_acc = %c0_f32 : f32) -> (f32) unroll { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %pair = func.call @qwen3_moe_q4k_q8_1_x4_paired_block_lane(%bounded_input_size, %weight, %weight_row_byte_base, %q8_input, %q8_row_byte_base, %q4_block, %bounded_cohort_lane) : (index, buffer, offset, buffer, offset, index, index) -> (f32) - %next = scalar.addf %block_acc, %pair : f32 - scf.yield %next : f32 - } - func.return %sum : f32 -} - -// Decodes the packed queue descriptor shared by expert-grouped projections. -// The producer stores a 7-bit expert ordinal, a 6-bit 32-row partition -// ordinal, and a 5-bit row count minus one. -func.def inline @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) { - %c1_i32 = scalar.constant 1 : i32 - %c5_i32 = scalar.constant 5 : i32 - %c7_i32 = scalar.constant 7 : i32 - %c13_i32 = scalar.constant 13 : i32 - %c31_i32 = scalar.constant 31 : i32 - %c63_i32 = scalar.constant 63 : i32 - %c127_i32 = scalar.constant 127 : i32 - %expert_i32 = scalar.andi %descriptor, %c127_i32 : i32 - %partition_shifted_i32 = scalar.shrui %descriptor, %c7_i32 : i32 - %partition_i32 = scalar.andi %partition_shifted_i32, %c63_i32 : i32 - %route_tile_base_i32 = scalar.shli %partition_i32, %c5_i32 : i32 - %row_count_shifted_i32 = scalar.shrui %descriptor, %c13_i32 : i32 - %row_count_minus_one_i32 = scalar.andi %row_count_shifted_i32, %c31_i32 : i32 - %partition_row_count_i32 = scalar.addi %row_count_minus_one_i32, %c1_i32 : i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index - %route_tile_base0 = index.cast %route_tile_base_i32 : i32 to index - %route_tile_base = index.assume %route_tile_base0 [range(%route_tile_base0, 0, 2016)] : index - %partition_row_count0 = index.cast %partition_row_count_i32 : i32 to index - %partition_row_count = index.assume %partition_row_count0 [range(%partition_row_count0, 1, 32)] : index - func.return %expert, %route_tile_base, %partition_row_count : index, index, index -} - -// Builds the transient expert table consumed by the grouped projection. The -// table packs [expert_count] route counts followed by -// [expert_count][token_count] compact assignment ordinals. Top-k routing -// selects each expert at most once per token, so token_count entries are -// sufficient for every expert even though there are route_count assignments -// per token. -kernel.def @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) { - %configured_expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - kernel.launch.config workgroups(%configured_expert_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 512)] : index - %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 512), eq(%expert_count, %configured_expert_count0)] : index, index - %expert0 = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %c0 = index.constant 0 : index - %workgroup_size = index.constant 256 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index - %expert, %table_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index - %assignment_table_byte_base = index.scale %table_expert_count, %c4_bytes : index, offset -> offset - %route_view = buffer.view %route_ids[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> - %count_view = buffer.view %expert_table[%c0_offset] : buffer -> view<[%table_expert_count]xi32> - %assignment_view = buffer.view %expert_table[%assignment_table_byte_base] : buffer -> view<[%table_expert_count]x[%bounded_token_count]xi32> - %expert_route_count = scf.for %block_base = [%c0 to %assignment_count step %workgroup_size](%matched_base = %c0_i32 : i32) -> (i32) { - %assignment = index.add %block_base, %lane : index - %in_range = index.cmp ult, %assignment, %assignment_count : index - %route_expert_i32 = scf.if %in_range -> (i32) { - %token0 = index.div %assignment, %configured_route_count : index - %route0 = index.rem %assignment, %configured_route_count : index - %token, %route_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - // The route stride is the physical row width and must contain every - // logical top-k route. - %route, %route_row_stride = index.assume %route0, %bounded_route_stride [lt(%route0, %bounded_route_stride)] : index, index - %loaded = view.load %route_view[%token, %route] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - scf.yield %loaded : i32 - } else { - %cn1_i32 = scalar.constant -1 : i32 - scf.yield %cn1_i32 : i32 - } - %route_expert0 = index.cast %route_expert_i32 : i32 to index - %route_expert = index.assume %route_expert0 [range(%route_expert0, -1, 511)] : index - %matches = index.cmp eq, %route_expert, %expert : index - %match_i32 = scf.if %matches -> (i32) { - scf.yield %c1_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %block_prefix = kernel.workgroup.scan %match_i32 {direction = forward, mode = exclusive} : i32 - %block_match_count_reduced = kernel.workgroup.reduce %match_i32 : i32 - %block_match_count = kernel.subgroup.broadcast.first %block_match_count_reduced : i32 - scf.if %matches { - %match_ordinal_i32 = scalar.addi %matched_base, %block_prefix : i32 - %match_ordinal0 = index.cast %match_ordinal_i32 : i32 to index - %match_ordinal = index.assume %match_ordinal0 [range(%match_ordinal0, 0, 2047)] : index - // Top-k route IDs are unique within a token, so one expert can own at - // most token_count assignments. - %bounded_match_ordinal, %table_token_count = index.assume %match_ordinal, %bounded_token_count [lt(%match_ordinal, %bounded_token_count)] : index, index - %assignment_i32 = index.cast %assignment : index to i32 - view.store %assignment_i32, %assignment_view[%expert, %bounded_match_ordinal] : i32, view<[%table_expert_count]x[%bounded_token_count]xi32> - } - %next_matched_base = scalar.addi %matched_base, %block_match_count : i32 - scf.yield %next_matched_base : i32 - } - %is_lane_zero = index.cmp eq, %lane, %c0 : index - scf.if %is_lane_zero { - view.store %expert_route_count, %count_view[%expert] : i32, view<[%table_expert_count]xi32> - } - kernel.return -} - -// Compacts expert assignment counts into exact 32-row projection partitions. -// -// One lane owns each expert. A workgroup scan assigns deterministic descriptor -// offsets, then each lane publishes an 18-bit descriptor containing the expert, -// 32-row partition ordinal, and tail row count. The consumer can launch a -// distribution-independent grid without serializing a concentrated expert -// inside one workgroup. -kernel.def @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) { - %c1 = index.constant 1 : index - %workgroup_size = index.constant 128 : index - kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %lane = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c0_i32 = scalar.constant 0 : i32 - %c1_i32 = scalar.constant 1 : i32 - %c7_i32 = scalar.constant 7 : i32 - %c13_i32 = scalar.constant 13 : i32 - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index - %count_view = buffer.view %expert_table[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> - %partition_count_view = buffer.view %partition_table[%c0_offset] : buffer -> view<1xi32> - %partition_descriptor_view = buffer.view %partition_table[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> - %has_expert = index.cmp ult, %lane, %bounded_expert_count : index - %expert_assignment_count_i32 = scf.if %has_expert -> (i32) { - %expert, %table_expert_count = index.assume %lane, %bounded_expert_count [lt(%lane, %bounded_expert_count)] : index, index - %loaded = view.load %count_view[%expert] : view<[%bounded_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_assignment_count0 = index.cast %expert_assignment_count_i32 : i32 to index - %expert_assignment_count = index.assume %expert_assignment_count0 [range(%expert_assignment_count0, 0, 2048)] : index - %rounded_expert_assignment_count = index.add %expert_assignment_count, %c31 : index - %expert_partition_count = index.div %rounded_expert_assignment_count, %c32 : index - %expert_partition_count_i32 = index.cast %expert_partition_count : index to i32 - %expert_partition_base_i32 = kernel.workgroup.scan %expert_partition_count_i32 {direction = forward, mode = exclusive} : i32 - %partition_count_i32 = kernel.workgroup.reduce %expert_partition_count_i32 : i32 - %partition_count = index.cast %partition_count_i32 : i32 to index - %bounded_partition_count, %table_partition_capacity = index.assume %partition_count, %maximum_partition_count [lt(%partition_count, %maximum_partition_count)] : index, index - %expert_partition_base0 = index.cast %expert_partition_base_i32 : i32 to index - %expert_partition_base = index.assume %expert_partition_base0 [range(%expert_partition_base0, 0, 639)] : index - scf.if %has_expert { - %expert, %table_expert_count = index.assume %lane, %bounded_expert_count [lt(%lane, %bounded_expert_count)] : index, index - %expert_i32 = index.cast %expert : index to i32 - scf.for %partition = [%c0 to %expert_partition_count step %c1] { - %descriptor_ordinal0 = index.add %expert_partition_base, %partition : index - %descriptor_ordinal, %descriptor_count = index.assume %descriptor_ordinal0, %bounded_partition_count [lt(%descriptor_ordinal0, %bounded_partition_count)] : index, index - %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index - %partition_remainder = index.rem %expert_assignment_count, %c32 : index - %has_partial_tail = index.cmp ne, %partition_remainder, %c0 : index - %partition_row_count = scf.if %has_partial_tail -> (index) { - %next_partition = index.add %partition, %c1 : index - %is_tail_partition = index.cmp eq, %next_partition, %expert_partition_count : index - %tail_row_count = scf.if %is_tail_partition -> (index) { - scf.yield %partition_remainder : index - } else { - scf.yield %c32 : index - } - scf.yield %tail_row_count : index - } else { - scf.yield %c32 : index - } - %partition_i32 = index.cast %partition : index to i32 - %partition_row_count_i32 = index.cast %partition_row_count : index to i32 - %packed_partition = scalar.shli %partition_i32, %c7_i32 : i32 - %partition_row_count_minus_one = scalar.subi %partition_row_count_i32, %c1_i32 : i32 - %packed_row_count = scalar.shli %partition_row_count_minus_one, %c13_i32 : i32 - %packed_expert_partition = scalar.ori %expert_i32, %packed_partition : i32 - %packed_descriptor = scalar.ori %packed_expert_partition, %packed_row_count : i32 - view.store %packed_descriptor, %partition_descriptor_view[%table_descriptor_ordinal] : i32, view<[%maximum_partition_count]xi32> - } - } - %is_lane_zero = index.cmp eq, %lane, %c0 : index - scf.if %is_lane_zero { - view.store %partition_count_i32, %partition_count_view[%c0] : i32, view<1xi32> - } - kernel.return -} - -// Small-token body for routed Q4_K gate/up projections. One wave computes one -// [token, route, output channel]. Callers own launch geometry and may append a -// producer epilogue after the F32 SwiGLU value is visible. -func.template device @qwen3_moe_routed_gate_up_swiglu_q4k_q8_body(%publish_output: i1, %token_count: index, %token: index, %route_count: index, %route: index, %route_stride: index, %expert_count: index, %output_size: index, %channel: index, %lane: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512)] : index - %bounded_route_count, %bounded_route_stride = index.assume %route_count, %route_stride [range(%route_count, 1, 8), range(%route_stride, 1, 128), le(%route_count, %route_stride)] : index, index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %bounded_token, %body_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %bounded_channel, %body_output_size = index.assume %channel, %bounded_output_size [lt(%channel, %bounded_output_size)] : index, index - %bounded_route, %body_route_count = index.assume %route, %bounded_route_count [lt(%route, %bounded_route_count)] : index, 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 - %c256 = index.constant 256 : index - %c1023 = index.constant 1023 : index - %c1024 = index.constant 1024 : index - %q4_block_bytes = index.constant 144 : index - %q4_block_bytes_offset = index.constant 144 : offset - %q8_group_bytes = index.constant 144 : offset - %q8_payload_byte_add = index.constant 16 : offset - %c1_byte = index.constant 1 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c0_i32 = scalar.constant 0 : i32 - %c0_offset = index.constant 0 : offset - %q4_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.mul %q4_block_count, %q4_block_bytes : index - %weight_expert_bytes = index.mul %bounded_output_size, %weight_row_bytes : index - %q8_group_count = index.div %input_size, %c128 : index - %q8_bytes_per_token = index.mul %q8_group_count, %q4_block_bytes : index - %padded_input_size = index.add %input_size, %c1023 : index - %iteration_count = index.div %padded_input_size, %c1024 : index - %q8_noalias, %route_ids_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %q8_input, %route_ids, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer - %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_stride]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_route_count]x[%bounded_output_size]xf32> - %route_index, %route_row_stride = index.assume %bounded_route, %bounded_route_stride [lt(%bounded_route, %bounded_route_stride)] : index, index - %expert_i32 = view.load %route_ids_view[%bounded_token, %route_index] : view<[%bounded_token_count]x[%bounded_route_stride]xi32> -> i32 - %expert0 = index.cast %expert_i32 : i32 to index - %expert = index.assume %expert0 [range(%expert0, 0, 127)] : index - %expert_byte_base = index.mul %expert, %weight_expert_bytes : index - %channel_byte_add = index.mul %bounded_channel, %weight_row_bytes : index - %row_byte_index = index.add %expert_byte_base, %channel_byte_add : index - %row_byte_base = index.scale %row_byte_index, %c1_byte : index, offset -> offset - %q8_token_byte_index = index.mul %bounded_token, %q8_bytes_per_token : index - %q8_token_byte_base = index.scale %q8_token_byte_index, %c1_byte : index, offset -> offset - %q4_group_pair0 = index.div %bounded_lane, %c2 : index - %q4_group_pair1 = index.rem %q4_group_pair0, %c4 : index - %q4_group_pair = index.assume %q4_group_pair1 [range(%q4_group_pair1, 0, 3)] : index - %q4_half0 = index.rem %bounded_lane, %c2 : index - %q4_half = index.assume %q4_half0 [range(%q4_half0, 0, 1)] : index - %lane_q4_block = index.div %bounded_lane, %c8 : index - %q8_group_in_block0 = index.div %q4_group_pair, %c2 : index - %q8_group_in_block = index.assume %q8_group_in_block0 [range(%q8_group_in_block0, 0, 1)] : index - %pair_in_q8_group0 = index.rem %q4_group_pair, %c2 : index - %pair_in_q8_group = index.assume %pair_in_q8_group0 [range(%pair_in_q8_group0, 0, 1)] : index - %q8_low_inner_block0 = index.mul %pair_in_q8_group, %c2 : index - %q8_low_inner_block = index.assume %q8_low_inner_block0 [range(%q8_low_inner_block0, 0, 2)] : index - %q8_high_inner_block0 = index.add %q8_low_inner_block, %c1 : index - %q8_high_inner_block = index.assume %q8_high_inner_block0 [range(%q8_high_inner_block0, 1, 3)] : index - %q8_half_word_add = index.mul %q4_half, %c4 : index - %q8_low_inner_word_base = index.mul %q8_low_inner_block, %c8 : index - %q8_low_word_index0 = index.add %q8_low_inner_word_base, %q8_half_word_add : index - %q8_low_word_index = index.assume %q8_low_word_index0 [range(%q8_low_word_index0, 0, 20)] : index - %q8_high_inner_word_base = index.mul %q8_high_inner_block, %c8 : index - %q8_high_word_index0 = index.add %q8_high_inner_word_base, %q8_half_word_add : index - %q8_high_word_index = index.assume %q8_high_word_index0 [range(%q8_high_word_index0, 8, 28)] : index - %q8_low_ds_index0 = index.mul %q8_low_inner_block, %c2 : index - %q8_low_ds_index = index.assume %q8_low_ds_index0 [range(%q8_low_ds_index0, 0, 4)] : index - %gate_acc, %up_acc = scf.for %iteration = [%c0 to %iteration_count step %c1](%gate_iter = %c0_f32 : f32, %up_iter = %c0_f32 : f32) -> (f32, f32) unroll { - %iteration_q4_block = index.mul %iteration, %c4 : index - %q4_block0 = index.add %iteration_q4_block, %lane_q4_block : index - %valid_q4_block = index.cmp ult, %q4_block0, %q4_block_count : index - %gate_contribution, %up_contribution = scf.if %valid_q4_block -> (f32, f32) { - %q4_block, %bounded_q4_block_count = index.assume %q4_block0, %q4_block_count [lt(%q4_block0, %q4_block_count)] : index, index - %q8_block_group_base = index.mul %q4_block, %c2 : index - %q8_group0 = index.add %q8_block_group_base, %q8_group_in_block : index - %q8_group, %bounded_q8_group_count = index.assume %q8_group0, %q8_group_count [lt(%q8_group0, %q8_group_count)] : index, index - %q8_group_byte_add = index.scale %q8_group, %q8_group_bytes : index, offset -> offset - %q8_group_byte_base = index.add %q8_token_byte_base, %q8_group_byte_add : offset - %q8_payload_byte_base = index.add %q8_group_byte_base, %q8_payload_byte_add : offset - %q8_ds_view = buffer.view %q8_noalias[%q8_group_byte_base] : buffer -> view<8xf16> - %q8_words_view = buffer.view %q8_noalias[%q8_payload_byte_base] : buffer -> view<32xi32> - %q8_ds = vector.load %q8_ds_view[%q8_low_ds_index] : view<8xf16> -> vector<4xf16> - %q8_low_d_f16 = vector.extract %q8_ds[0] : vector<4xf16> -> f16 - %q8_low_s_f16 = vector.extract %q8_ds[1] : vector<4xf16> -> f16 - %q8_high_d_f16 = vector.extract %q8_ds[2] : vector<4xf16> -> f16 - %q8_high_s_f16 = vector.extract %q8_ds[3] : vector<4xf16> -> f16 - %q8_low_d = scalar.extf %q8_low_d_f16 : f16 to f32 - %q8_low_s = scalar.extf %q8_low_s_f16 : f16 to f32 - %q8_high_d = scalar.extf %q8_high_d_f16 : f16 to f32 - %q8_high_s = scalar.extf %q8_high_s_f16 : f16 to f32 - %q8_low_words = vector.load %q8_words_view[%q8_low_word_index] : view<32xi32> -> vector<4xi32> - %q8_high_words = vector.load %q8_words_view[%q8_high_word_index] : view<32xi32> -> vector<4xi32> - %q8_low_values = vector.bitcast %q8_low_words : vector<4xi32> to vector<16xi8> - %q8_high_values = vector.bitcast %q8_high_words : vector<4xi32> to vector<16xi8> - %q4_block_byte_add = index.scale %q4_block, %q4_block_bytes_offset : index, offset -> offset - %q4_block_byte_base = index.add %row_byte_base, %q4_block_byte_add : offset - %gate_header_view = buffer.view %gate_noalias[%q4_block_byte_base] : buffer -> view<4xi32> - %up_header_view = buffer.view %up_noalias[%q4_block_byte_base] : buffer -> view<4xi32> - %gate_header_words = vector.load %gate_header_view[0] : view<4xi32> -> vector<4xi32> - %up_header_words = vector.load %up_header_view[0] : view<4xi32> -> vector<4xi32> - %gate_q4_low, %gate_low_d_scale, %gate_low_dmin_scale, %gate_q4_high, %gate_high_d_scale, %gate_high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%gate_noalias, %row_byte_base, %q4_block, %q4_group_pair, %q4_half, %gate_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) - %up_q4_low, %up_low_d_scale, %up_low_dmin_scale, %up_q4_high, %up_high_d_scale, %up_high_dmin_scale = func.call @qwen3_moe_q4k_chunk_pair_global(%up_noalias, %row_byte_base, %q4_block, %q4_group_pair, %q4_half, %up_header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) - %gate_low = func.call @qwen3_moe_q4k_q8_1_dot(%gate_q4_low, %gate_low_d_scale, %gate_low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %gate_high = func.call @qwen3_moe_q4k_q8_1_dot(%gate_q4_high, %gate_high_d_scale, %gate_high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %up_low = func.call @qwen3_moe_q4k_q8_1_dot(%up_q4_low, %up_low_d_scale, %up_low_dmin_scale, %q8_low_values, %q8_low_d, %q8_low_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %up_high = func.call @qwen3_moe_q4k_q8_1_dot(%up_q4_high, %up_high_d_scale, %up_high_dmin_scale, %q8_high_values, %q8_high_d, %q8_high_s) : (vector<16xi8>, f32, f32, vector<16xi8>, f32, f32) -> (f32) - %gate_pair = scalar.addf %gate_low, %gate_high : f32 - %up_pair = scalar.addf %up_low, %up_high : f32 - scf.yield %gate_pair, %up_pair : f32, f32 - } else { - scf.yield %c0_f32, %c0_f32 : f32, f32 - } - %gate_next = scalar.addf %gate_iter, %gate_contribution : f32 - %up_next = scalar.addf %up_iter, %up_contribution : f32 - scf.yield %gate_next, %up_next : f32, f32 - } - %gate_dot = kernel.subgroup.reduce %gate_acc : f32 - %up_dot = kernel.subgroup.reduce %up_acc : f32 - %lane_i32 = index.cast %bounded_lane : index to i32 - %is_lane_zero = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - %writes_output = scalar.andi %publish_output, %is_lane_zero : i1 - scf.if %writes_output { - %gate_silu = scalar.siluf %gate_dot : f32 - %result = scalar.mulf %gate_silu, %up_dot : f32 - view.store %result, %output_view[%bounded_token, %bounded_route, %bounded_channel] : f32, view<[%bounded_token_count]x[%bounded_route_count]x[%bounded_output_size]xf32> - } - func.return -} - -// Small-token schedule for routed Q4_K gate/up projections. Each workgroup -// packs four independent channel waves for one [token, route] pair. Eight -// lanes consume both nibbles of the packed codes for one 256-element block, so -// each wave advances through K in 1024-element stripes. Gate and up share each -// paired Q8_1 load before independent raw-weight dot products and a fused -// SwiGLU epilogue. -kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) { - %configured_route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %unit = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %workgroup_size = index.constant 128 : index - %padded_output_size = index.add %configured_output_size, %c3 : index - %channel_workgroup_count = index.div %padded_output_size, %c4 : index - kernel.launch.config workgroups(%channel_workgroup_count, %configured_route_count, %token_count) workgroup_size(%workgroup_size, %unit, %unit) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { - %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512)] : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 128)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %token0 = kernel.workgroup.id : index - %channel_workgroup = kernel.workgroup.id : index - %route0 = kernel.workgroup.id : index - %subgroup = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c4 = index.constant 4 : index - %valid_token = index.cmp ult, %token0, %bounded_token_count : index - %safe_token0 = scf.select %valid_token, %token0, %c0 : index - %safe_token, %body_token_count = index.assume %safe_token0, %bounded_token_count [lt(%safe_token0, %bounded_token_count)] : index, index - %channel_base = index.mul %channel_workgroup, %c4 : index - %channel0 = index.add %channel_base, %subgroup : index - %valid_channel = index.cmp ult, %channel0, %bounded_output_size : index - %safe_channel0 = scf.select %valid_channel, %channel0, %c0 : index - %safe_channel, %body_output_size = index.assume %safe_channel0, %bounded_output_size [lt(%safe_channel0, %bounded_output_size)] : index, index - %route, %body_route_count = index.assume %route0, %bounded_route_count [lt(%route0, %bounded_route_count)] : index, index - %publishes_output = scalar.andi %valid_token, %valid_channel : i1 - func.apply(%publishes_output, %body_token_count, %safe_token, %body_route_count, %route, %bounded_route_stride, %bounded_expert_count, %body_output_size, %safe_channel, %lane, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - kernel.return -} - -// Packs one Q8_1 x4 physical group with one selected subgroup. Each lane owns -// four adjacent values and contributes its maximum and sum to one of four -// eight-lane logical blocks. One-hot vectors let ordinary subgroup reductions -// compute all four block aggregates without LDS or workgroup barriers. -func.template device @qwen3_moe_routed_gate_up_quantize_q8_1_x4_subgroup_body(%publish_output: i1, %group_count0: index, %group0: index, %input: buffer, %output: buffer) { - %group_count, %group = index.assume %group_count0, %group0 [range(%group_count0, 1, 524288), lt(%group0, %group_count0)] : index, index - %lane0 = kernel.subgroup.lane.id : index - %lane = index.assume %lane0 [range(%lane0, 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 - %group_bytes = index.constant 144 : offset - %payload_byte_add = index.constant 16 : offset - %c0_f32 = scalar.constant 0.0 : f32 - %c1_f32 = scalar.constant 1.0 : f32 - %c127_f32 = scalar.constant 127.0 : f32 - %c0_f32x4 = vector.constant 0.0 : vector<4xf32> - %c0_offset = index.constant 0 : offset - %launched_element_count = index.mul %group_count, %c128 : index - %group_element_base = index.mul %group, %c128 : index - %lane_element_add = index.mul %lane, %c4 : index - %input_index = index.add %group_element_base, %lane_element_add : index - %input_noalias, %output_noalias = buffer.assume.noalias %input, %output : buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launched_element_count]xf32> - %input_values = vector.load %input_view[%input_index] : view<[%launched_element_count]xf32> -> vector<4xf32> - %absolute_values = vector.absf %input_values : vector<4xf32> - %thread_maximum = vector.reduce %absolute_values, %c0_f32 : vector<4xf32>, f32 - %block0 = index.div %lane, %c8 : index - %block = index.assume %block0 [range(%block0, 0, 3)] : index - %block_maxima = vector.insert %thread_maximum into %c0_f32x4[%block] : f32, vector<4xf32> - %reduced_maxima = kernel.subgroup.reduce %block_maxima : vector<4xf32> - %amax = vector.extract %reduced_maxima[%block] : vector<4xf32> -> f32 - %d = scalar.divf %amax, %c127_f32 : f32 - %d_nonzero = scalar.cmpf one, %d, %c0_f32 : f32 - %d_inverse = scf.if %d_nonzero -> (f32) { - %inverse = scalar.divf %c1_f32, %d : f32 - scf.yield %inverse : f32 - } else { - scf.yield %c0_f32 : f32 - } - %d_inverse_vector = vector.splat %d_inverse : vector<4xf32> - %scaled_values = vector.mulf %input_values, %d_inverse_vector : vector<4xf32> - %rounded_values = vector.roundf %scaled_values : vector<4xf32> - %quantized_values = vector.fptosi %rounded_values : vector<4xf32> to vector<4xi8> - %packed_word = vector.bitcast %quantized_values : vector<4xi8> to vector<1xi32> - %group_byte_offset = index.scale %group, %group_bytes : index, offset -> offset - %payload_byte_offset = index.add %group_byte_offset, %payload_byte_add : offset - %group_ds = buffer.view %output_noalias[%group_byte_offset] : buffer -> view<8xf16> - %group_qs = buffer.view %output_noalias[%payload_byte_offset] : buffer -> view<32xi32> - scf.if %publish_output { - vector.store %packed_word, %group_qs[%lane] : vector<1xi32>, view<32xi32> - } - %thread_sum = vector.reduce %rounded_values, %c0_f32 : vector<4xf32>, f32 - %block_sums0 = vector.insert %thread_sum into %c0_f32x4[%block] : f32, vector<4xf32> - %reduced_sums = kernel.subgroup.reduce %block_sums0 : vector<4xf32> - %quantized_sum = vector.extract %reduced_sums[%block] : vector<4xf32> -> f32 - %s = scalar.mulf %quantized_sum, %d : f32 - %word_in_block = index.rem %lane, %c8 : index - %is_block_leader = index.cmp eq, %word_in_block, %c0 : index - %publishes_metadata = scalar.andi %publish_output, %is_block_leader : i1 - scf.if %publishes_metadata { - %d_f16 = scalar.fptrunc %d : f32 to f16 - %s_f16 = scalar.fptrunc %s : f32 to f16 - %ds_index = index.mul %block, %c2 : index - %s_index = index.add %ds_index, %c1 : index - view.store %d_f16, %group_ds[%ds_index] : f16, view<8xf16> - view.store %s_f16, %group_ds[%s_index] : f16, view<8xf16> - } - func.return -} - -// Decode producer that publishes both the ordinary F32 SwiGLU rows and their -// packed Q8_1 x4 representation. Four independent channel waves share one -// workgroup and publish one completion arrival. The last workgroup to complete -// a 128-channel physical group packs all four logical Q8_1 blocks with one -// selected subgroup, then resets the counter before returning. -kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) { - %configured_route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %unit = index.constant 1 : index - %c3 = index.constant 3 : index - %c4 = index.constant 4 : index - %workgroup_size = index.constant 128 : index - %padded_output_size = index.add %configured_output_size, %c3 : index - %channel_workgroup_count = index.div %padded_output_size, %c4 : index - kernel.launch.config workgroups(%channel_workgroup_count, %configured_route_count, %unit) workgroup_size(%workgroup_size, %unit, %unit) : index -} launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer, %completion_counters: buffer, %next_q8_output: buffer) { - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 1)] : index - %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_route_stride = index.assume %route_stride [range(%route_stride, 1, 128)] : index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 128, 4096), mul(%output_size, 128), eq(%output_size, %configured_output_size0)] : index, index - %q8_input_noalias, %route_ids_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias, %completion_counters_noalias, %next_q8_output_noalias = buffer.assume.noalias %q8_input, %route_ids, %gate_weight, %up_weight, %output, %completion_counters, %next_q8_output : buffer, buffer, buffer, buffer, buffer, buffer, buffer - %publishes_swiglu = scalar.constant true : i1 - %body_token = index.constant 0 : index - %channel_workgroup = kernel.workgroup.id : index - %route = kernel.workgroup.id : index - %token = kernel.workgroup.id : index - %subgroup = kernel.subgroup.id : index - %lane = kernel.subgroup.lane.id : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c128 = index.constant 128 : index - %c0_i32 = scalar.constant 0 : i32 - %c4_i32 = scalar.constant 4 : i32 - %c0_offset = index.constant 0 : offset - %counter_scratch_bytes = index.constant 4 : offset - %channel_base = index.mul %channel_workgroup, %c4 : index - %channel = index.add %channel_base, %subgroup : index - func.apply(%publishes_swiglu, %bounded_token_count, %body_token, %bounded_route_count, %route, %bounded_route_stride, %bounded_expert_count, %bounded_output_size, %channel, %lane, %q8_input_noalias, %route_ids_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias) : (i1, index, index, index, index, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) - %physical_group_count = index.div %bounded_output_size, %c128 : index - %row_count = index.mul %bounded_token_count, %bounded_route_count : index - %completion_counter_count = index.mul %row_count, %physical_group_count : index - %token_row_base = index.mul %token, %bounded_route_count : index - %row = index.add %token_row_base, %route : index - %row_group_base = index.mul %row, %physical_group_count : index - %group_in_row = index.div %channel_base, %c128 : index - %counter_index0 = index.add %row_group_base, %group_in_row : index - %counter_index, %bounded_completion_counter_count = index.assume %counter_index0, %completion_counter_count [lt(%counter_index0, %completion_counter_count)] : index, index - %completion_counters_aligned = buffer.assume.alignment %completion_counters_noalias {minimum_alignment = 16} : buffer - %completion_counters_view = buffer.view %completion_counters_aligned[%c0_offset] : buffer -> view<[%bounded_completion_counter_count]xi32> - %counter_scratch = buffer.alloca align(4) %counter_scratch_bytes : buffer - %counter_scratch_view = buffer.view %counter_scratch[%c0_offset] : buffer -> view<1xi32> - %is_arrival_lane = index.cmp eq, %workitem, %c0 : index - // Publish every producer lane's SwiGLU store before the leader advances one - // workgroup arrival. The last arrival then acquires the physical group. - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_lane { - %old_counter = view.atomic.rmw %c4_i32, %completion_counters_view[%counter_index] {ordering = acq_rel, scope = device} : i32, view<[%bounded_completion_counter_count]xi32> -> i32 - view.store %old_counter, %counter_scratch_view[%c0] : i32, view<1xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %old_counter = view.load %counter_scratch_view[%c0] : view<1xi32> -> i32 - %group_size_i32 = index.cast %c128 : index to i32 - %last_arrival_i32 = scalar.subi %group_size_i32, %c4_i32 : i32 - %negative_group_size_i32 = scalar.subi %c0_i32, %group_size_i32 : i32 - %is_last_arrival = scalar.cmpi eq, %old_counter, %last_arrival_i32 : i32 - scf.if %is_last_arrival { - kernel.barrier scope(workgroup) ordering(acquire) - %publish_output = scalar.constant true : i1 - %is_quantize_subgroup = index.cmp eq, %subgroup, %c0 : index - scf.if %is_quantize_subgroup { - func.apply(%publish_output, %bounded_completion_counter_count, %counter_index, %output_noalias, %next_q8_output_noalias) : (i1, index, index, buffer, buffer) - } - kernel.barrier scope(workgroup) ordering(release) - scf.if %is_arrival_lane { - view.atomic.reduce %negative_group_size_i32, %completion_counters_view[%counter_index] {ordering = release, scope = device} : i32, view<[%bounded_completion_counter_count]xi32> - } - } - kernel.return -} - -// Large-token schedule for routed Q4_K gate/up projections. Each workgroup -// owns one expert, 32 output channels, and up to 32 selected rows. It -// reads those rows from the compact expert table, stages one K=256 slice of -// both raw Q4_K projections and Q8_1 activations in LDS, and reuses each -// decoded weight chunk across four rows before a fused SwiGLU epilogue. -kernel.def @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped(%token_count: index, %route_count: index, %expert_count: index, %output_size: index) { - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %workgroup_size = index.constant 256 : index - %padded_output_size = index.add %configured_output_size, %c31 : index - %output_tiles = index.div %padded_output_size, %c32 : index - %padded_token_count = index.add %token_capacity, %c31 : index - %route_tiles = index.div %padded_token_count, %c32 : index - %route_partitions = index.min %route_tiles, %c4 : index - kernel.launch.config workgroups(%output_tiles, %route_partitions, %configured_expert_count) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %route_count: index, %expert_count: index, %output_size: index, %q8_input: buffer, %expert_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %token_capacity = config.get @qwen3_moe.workload.token_capacity : index - %configured_route_count0 = config.get @qwen3_moe.routed_gate_up.route_count : index - %configured_expert_count0 = config.get @qwen3_moe.routed_gate_up.expert_count : index - %configured_output_size0 = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 512), le(%token_count, %token_capacity)] : index - %bounded_route_count, %configured_route_count = index.assume %route_count, %configured_route_count0 [range(%route_count, 1, 8), eq(%route_count, %configured_route_count0)] : index, index - %bounded_expert_count, %configured_expert_count = index.assume %expert_count, %configured_expert_count0 [range(%expert_count, 1, 128), eq(%expert_count, %configured_expert_count0)] : index, index - %bounded_output_size, %configured_output_size = index.assume %output_size, %configured_output_size0 [range(%output_size, 1, 4096), eq(%output_size, %configured_output_size0)] : index, index - %channel_tile = kernel.workgroup.id : index - %route_tile = kernel.workgroup.id : index - %expert0 = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c32 = index.constant 32 : index - %c36 = index.constant 36 : index - %c72 = index.constant 72 : index - %c128 = index.constant 128 : index - %c256 = index.constant 256 : index - %gate_stage_word_count = index.constant 1152 : index - %q8_stage_word_count = index.constant 2304 : index - %c0_i32 = scalar.constant 0 : i32 - %c0_f32_rows = vector.constant 0.0 : vector<4xf32> - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %q8_row_bytes = index.constant 288 : offset - %row_id_bytes = index.constant 128 : offset - %weight_stage_bytes = index.constant 4608 : offset - %q8_stage_bytes = index.constant 9216 : offset - %assignment_count = index.mul %bounded_token_count, %bounded_route_count : index - %expert, %table_expert_count = index.assume %expert0, %bounded_expert_count [lt(%expert0, %bounded_expert_count)] : index, index - %assignment_table_byte_base = index.scale %table_expert_count, %c4_bytes : index, offset -> offset - %q4_block_count = index.div %input_size, %c256 : index - %output_route_count = index.mul %bounded_token_count, %bounded_route_count : index - %q8_noalias, %expert_table_noalias, %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %q8_input, %expert_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer - %q8_words = buffer.view %q8_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%q4_block_count]x72xi32> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%table_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%table_expert_count]x[%bounded_token_count]xi32> - %gate_words = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> - %up_words = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_route_count]x[%bounded_output_size]xf32> - %row_ids = buffer.alloca align(16) %row_id_bytes : buffer - %gate_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %up_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %q8_stage = buffer.alloca align(16) %q8_stage_bytes : buffer - %row_ids_view = buffer.view %row_ids[%c0_offset] : buffer -> view<32xi32> - %gate_stage_words = buffer.view %gate_stage[%c0_offset] : buffer -> view<1152xi32> - %up_stage_words = buffer.view %up_stage[%c0_offset] : buffer -> view<1152xi32> - %q8_stage_words = buffer.view %q8_stage[%c0_offset] : buffer -> view<2304xi32> - %channel_base = index.mul %channel_tile, %c32 : index - %channel0 = index.rem %lane, %c32 : index - %channel = index.add %channel_base, %channel0 : index - %row_base0 = index.div %lane, %c32 : index - %row_base = index.assume %row_base0 [range(%row_base0, 0, 7)] : index - %initial_tile_base = index.mul %route_tile, %c32 : index - // Four interleaved route partitions cover the first 128 rows. Smaller - // token counts execute at most one iteration; larger expert populations - // continue in 128-row strides without increasing the launch grid. - %route_partition_step = index.constant 128 : index - %is_lane_zero = index.cmp eq, %lane, %c0 : index - %lane_expert_route_count = scf.if %is_lane_zero -> (i32) { - %loaded = view.load %count_view[%expert] : view<[%table_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 - %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 - %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index - %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 4096)] : index - scf.for %tile_base = [%initial_tile_base to %expert_route_count step %route_partition_step] { - %tile_base_i32 = index.cast %tile_base : index to i32 - %remaining_rows_i32 = scalar.subi %expert_route_count_i32, %tile_base_i32 : i32 - %remaining_rows0 = index.cast %remaining_rows_i32 : i32 to index - %remaining_rows = index.assume %remaining_rows0 [range(%remaining_rows0, 1, 4096)] : index - %has_full_tile = index.cmp uge, %remaining_rows, %c32 : index - %tile_row_count = scf.if %has_full_tile -> (index) { - scf.yield %c32 : index - } else { - scf.yield %remaining_rows : index - } - %loads_row = index.cmp ult, %lane, %tile_row_count : index - scf.if %loads_row { - %local_row = index.assume %lane [range(%lane, 0, 31)] : index - %expert_assignment_ordinal0 = index.add %tile_base, %local_row : index - %expert_assignment_ordinal, %table_token_count = index.assume %expert_assignment_ordinal0, %bounded_token_count [lt(%expert_assignment_ordinal0, %bounded_token_count)] : index, index - %assignment_i32 = view.load %assignment_view[%expert, %expert_assignment_ordinal] : view<[%table_expert_count]x[%bounded_token_count]xi32> -> i32 - view.store %assignment_i32, %row_ids_view[%local_row] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %gate_acc, %up_acc = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%gate_block_acc = %c0_f32_rows : vector<4xf32>, %up_block_acc = %c0_f32_rows : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) { - %bounded_q4_block, %body_q4_block_count = index.assume %q4_block, %q4_block_count [lt(%q4_block, %q4_block_count)] : index, index - scf.for %stage_word = [%lane to %gate_stage_word_count step %c256] { - %local_channel = index.div %stage_word, %c36 : index - %word_in_block0 = index.rem %stage_word, %c36 : index - %word_in_block = index.assume %word_in_block0 [range(%word_in_block0, 0, 35)] : index - %global_channel = index.add %channel_base, %local_channel : index - %valid_channel = index.cmp ult, %global_channel, %bounded_output_size : index - %gate_word, %up_word = scf.if %valid_channel -> (i32, i32) { - %bounded_global_channel, %table_output_size = index.assume %global_channel, %bounded_output_size [lt(%global_channel, %bounded_output_size)] : index, index - %gate_loaded = view.load %gate_words[%expert, %bounded_global_channel, %bounded_q4_block, %word_in_block] : view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> -> i32 - %up_loaded = view.load %up_words[%expert, %bounded_global_channel, %bounded_q4_block, %word_in_block] : view<[%bounded_expert_count]x[%bounded_output_size]x[%q4_block_count]x36xi32> -> i32 - scf.yield %gate_loaded, %up_loaded : i32, i32 - } else { - scf.yield %c0_i32, %c0_i32 : i32, i32 - } - view.store %gate_word, %gate_stage_words[%stage_word] : i32, view<1152xi32> - view.store %up_word, %up_stage_words[%stage_word] : i32, view<1152xi32> - } - scf.for %stage_word = [%lane to %q8_stage_word_count step %c256] { - %local_row = index.div %stage_word, %c72 : index - %word_in_row0 = index.rem %stage_word, %c72 : index - %word_in_row = index.assume %word_in_row0 [range(%word_in_row0, 0, 71)] : index - %valid_row = index.cmp ult, %local_row, %tile_row_count : index - %q8_word = scf.if %valid_row -> (i32) { - %bounded_local_row = index.assume %local_row [range(%local_row, 0, 31)] : index - %assignment_i32 = view.load %row_ids_view[%bounded_local_row] : view<32xi32> -> i32 - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 4095)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %token0 = index.div %bounded_assignment, %configured_route_count : index - %token, %table_token_count = index.assume %token0, %bounded_token_count [lt(%token0, %bounded_token_count)] : index, index - %loaded = view.load %q8_words[%token, %bounded_q4_block, %word_in_row] : view<[%bounded_token_count]x[%q4_block_count]x72xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - view.store %q8_word, %q8_stage_words[%stage_word] : i32, view<2304xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %channel_byte_base = index.scale %channel0, %q4_block_bytes : index, offset -> offset - %after_groups_gate, %after_groups_up = scf.for %q4_group = [%c0 to %c8 step %c1](%gate_group_acc = %gate_block_acc : vector<4xf32>, %up_group_acc = %up_block_acc : vector<4xf32>) -> (vector<4xf32>, vector<4xf32>) unroll { - %row1 = index.add %row_base, %c8 : index - %row2_add = index.constant 16 : index - %row2 = index.add %row_base, %row2_add : index - %row3_add = index.constant 24 : index - %row3 = index.add %row_base, %row3_add : index - %row0_byte_base = index.scale %row_base, %q8_row_bytes : index, offset -> offset - %row1_byte_base = index.scale %row1, %q8_row_bytes : index, offset -> offset - %row2_byte_base = index.scale %row2, %q8_row_bytes : index, offset -> offset - %row3_byte_base = index.scale %row3, %q8_row_bytes : index, offset -> offset - %q8_values0, %q8_d0, %q8_s0 = func.call @ggml_q8_1_x4_block(%q8_stage, %row0_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_values1, %q8_d1, %q8_s1 = func.call @ggml_q8_1_x4_block(%q8_stage, %row1_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_values2, %q8_d2, %q8_s2 = func.call @ggml_q8_1_x4_block(%q8_stage, %row2_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_values3, %q8_d3, %q8_s3 = func.call @ggml_q8_1_x4_block(%q8_stage, %row3_byte_base, %q4_group) : (buffer, offset, index) -> (vector<32xi8>, f32, f32) - %q8_low0 = vector.slice %q8_values0[0] : vector<32xi8> -> vector<16xi8> - %q8_low1 = vector.slice %q8_values1[0] : vector<32xi8> -> vector<16xi8> - %q8_low2 = vector.slice %q8_values2[0] : vector<32xi8> -> vector<16xi8> - %q8_low3 = vector.slice %q8_values3[0] : vector<32xi8> -> vector<16xi8> - %q8_high0 = vector.slice %q8_values0[16] : vector<32xi8> -> vector<16xi8> - %q8_high1 = vector.slice %q8_values1[16] : vector<32xi8> -> vector<16xi8> - %q8_high2 = vector.slice %q8_values2[16] : vector<32xi8> -> vector<16xi8> - %q8_high3 = vector.slice %q8_values3[16] : vector<32xi8> -> vector<16xi8> - %q8_rows_low0 = vector.constant 0 : vector<4x16xi8> - %q8_rows_low1 = vector.insert %q8_low0 into %q8_rows_low0[0] : vector<16xi8>, vector<4x16xi8> - %q8_rows_low2 = vector.insert %q8_low1 into %q8_rows_low1[1] : vector<16xi8>, vector<4x16xi8> - %q8_rows_low3 = vector.insert %q8_low2 into %q8_rows_low2[2] : vector<16xi8>, vector<4x16xi8> - %q8_rows_low = vector.insert %q8_low3 into %q8_rows_low3[3] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high0 = vector.constant 0 : vector<4x16xi8> - %q8_rows_high1 = vector.insert %q8_high0 into %q8_rows_high0[0] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high2 = vector.insert %q8_high1 into %q8_rows_high1[1] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high3 = vector.insert %q8_high2 into %q8_rows_high2[2] : vector<16xi8>, vector<4x16xi8> - %q8_rows_high = vector.insert %q8_high3 into %q8_rows_high3[3] : vector<16xi8>, vector<4x16xi8> - %q8_d = vector.from_elements %q8_d0, %q8_d1, %q8_d2, %q8_d3 : vector<4xf32> - %q8_s = vector.from_elements %q8_s0, %q8_s1, %q8_s2, %q8_s3 : vector<4xf32> - %gate_q4_low, %gate_d_scale_low, %gate_dmin_scale_low = func.call @qwen3_moe_q4k_chunk_local(%gate_stage, %channel_byte_base, %c0, %q4_group, %c0) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %gate_q4_high, %gate_d_scale_high, %gate_dmin_scale_high = func.call @qwen3_moe_q4k_chunk_local(%gate_stage, %channel_byte_base, %c0, %q4_group, %c1) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %up_q4_low, %up_d_scale_low, %up_dmin_scale_low = func.call @qwen3_moe_q4k_chunk_local(%up_stage, %channel_byte_base, %c0, %q4_group, %c0) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %up_q4_high, %up_d_scale_high, %up_dmin_scale_high = func.call @qwen3_moe_q4k_chunk_local(%up_stage, %channel_byte_base, %c0, %q4_group, %c1) : (buffer, offset, index, index, index) -> (vector<16xi8>, f32, f32) - %gate_low = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%gate_q4_low, %gate_d_scale_low, %gate_dmin_scale_low, %q8_rows_low, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %gate_high = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%gate_q4_high, %gate_d_scale_high, %gate_dmin_scale_high, %q8_rows_high, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %up_low = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%up_q4_low, %up_d_scale_low, %up_dmin_scale_low, %q8_rows_low, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %up_high = func.call @qwen3_moe_q4k_q8_1_dot4_rows(%up_q4_high, %up_d_scale_high, %up_dmin_scale_high, %q8_rows_high, %q8_d, %q8_s) : (vector<16xi8>, f32, f32, vector<4x16xi8>, vector<4xf32>, vector<4xf32>) -> (vector<4xf32>) - %gate_pair = vector.addf %gate_low, %gate_high : vector<4xf32> - %up_pair = vector.addf %up_low, %up_high : vector<4xf32> - %gate_next = vector.addf %gate_group_acc, %gate_pair : vector<4xf32> - %up_next = vector.addf %up_group_acc, %up_pair : vector<4xf32> - scf.yield %gate_next, %up_next : vector<4xf32>, vector<4xf32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %after_groups_gate, %after_groups_up : vector<4xf32>, vector<4xf32> - } - %gate_silu = vector.siluf %gate_acc : vector<4xf32> - %result_rows = vector.mulf %gate_silu, %up_acc : vector<4xf32> - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - scf.for %row_variant = [%c0 to %c4 step %c1] unroll { - %row_add = index.mul %row_variant, %c8 : index - %local_row = index.add %row_base, %row_add : index - %valid_row = index.cmp ult, %local_row, %tile_row_count : index - %writes_output = scalar.andi %valid_channel, %valid_row : i1 - scf.if %writes_output { - %bounded_local_row = index.assume %local_row [range(%local_row, 0, 31)] : index - %assignment_i32 = view.load %row_ids_view[%bounded_local_row] : view<32xi32> -> i32 - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 4095)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %result = vector.extract %result_rows[%row_variant] : vector<4xf32> -> f32 - view.store %result, %output_view[%bounded_assignment, %channel] : f32, view<[%output_route_count]x[%bounded_output_size]xf32> - } - } - } - kernel.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_nonzero_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(1) : index - %expert_count = check.literal value(1) : index - %output_size = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> - %q8_input = check.generate.fill value(0) : tensor<1x2304xi8> - %route_ids = check.generate.fill value(0) : tensor<1xi32> - %gate_weight = check.generate.fill value(85) : tensor<8x144xi8> - %up_weight = check.generate.fill value(-86) : tensor<8x144xi8> - %output = check.generate.fill value(0.0) : tensor<1xf32> - %expected = check.generate.fill value(-9024645.0) : tensor<1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<1x2304xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<1x2304xi8>, tensor<1xi32>, tensor<8x144xi8>, tensor<8x144xi8>, tensor<1xf32>) - check.expect.close actual(%output) expected(%expected) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<1xf32> - check.return -} - -// Exact decode topology with compact route IDs and every physical output -// group populated. Two fused invocations prove counter reuse in addition to -// comparing the F32 and packed outputs with the ordinary composition. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_differential_case { - %token_count = check.literal value(1) : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %routed_row_count = check.literal value(8) : index - %input = check.generate.fill value(0.00390625) : tensor<1x2048xf32> - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<128x768x8x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<128x768x8x144xi8> - %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> - %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> - %actual_output0 = check.generate.fill value(1.0) : tensor<8x768xf32> - %actual_q8_0 = check.generate.fill value(1) : tensor<8x864xi8> - %actual_output1 = check.generate.fill value(2.0) : tensor<8x768xf32> - %actual_q8_1 = check.generate.fill value(2) : tensor<8x864xi8> - %completion_counters = check.generate.fill value(0) : tensor<48xi32> - %expected_counters = check.generate.fill value(0) : tensor<48xi32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<1x2048xf32>, tensor<2304xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %output_size](%routed_row_count, %output_size, %expected_output, %expected_q8) : [index, index](index, index, tensor<8x768xf32>, tensor<8x864xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %actual_output0, %completion_counters, %actual_q8_0) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %actual_output1, %completion_counters, %actual_q8_1) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) - check.expect.close actual(%actual_output0) expected(%expected_output) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<8x768xf32> - check.expect.close actual(%actual_output1) expected(%expected_output) atol(16.0) rtol(9.9999999999999995e-07) nan(same) : tensor<8x768xf32> - check.expect.equal actual(%actual_q8_0) expected(%expected_q8) : tensor<8x864xi8> - check.expect.equal actual(%actual_q8_1) expected(%expected_q8) : tensor<8x864xi8> - check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<48xi32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_differential_case { - %token_count = check.literal value(64) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(32) : index - %input = check.generate.fill value(0.00390625) : tensor<64x512xf32> - %q8_input = check.generate.fill value(0) : tensor<64x576xi8> - // The physical row retains four argsort entries while the logical top-k view - // selects experts 0 and 1. Both experts span two 32-row grouped tiles. - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<64x4xi32> - %expert_table = check.generate.fill value(-1) : tensor<260xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<4x32x2x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<4x32x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<64x2x32xf32> - %actual = check.generate.fill value(1.0) : tensor<64x2x32xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<64x512xf32>, tensor<64x576xi8>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<64x4xi32>, tensor<260xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<64x576xi8>, tensor<64x4xi32>, tensor<4x32x2x144xi8>, tensor<4x32x2x144xi8>, tensor<64x2x32xf32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<64x576xi8>, tensor<260xi32>, tensor<4x32x2x144xi8>, tensor<4x32x2x144xi8>, tensor<64x2x32xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<64x2x32xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_tail_case { - %token_count = check.literal value(37) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(3) : index - %route_stride = check.literal value(5) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %input = check.generate.fill value(0.00390625) : tensor<37x512xf32> - %q8_input = check.generate.fill value(0) : tensor<37x576xi8> - // Three unique selected experts rotate through a five-entry physical row. - // This leaves the second route tile empty and the second channel tile with - // only one live output channel. - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<37x5xi32> - %expert_table = check.generate.fill value(-1) : tensor<152xi32> - %gate_weight = check.generate.iota offset(-72) step(1) period(144) : tensor<4x33x2x144xi8> - %up_weight = check.generate.iota offset(-71) step(1) period(144) : tensor<4x33x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<37x3x33xf32> - %actual = check.generate.fill value(1.0) : tensor<37x3x33xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<37x512xf32>, tensor<37x576xi8>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<37x5xi32>, tensor<152xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<37x576xi8>, tensor<37x5xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<37x3x33xf32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<37x576xi8>, tensor<152xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<37x3x33xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<37x3x33xf32> - check.return -} - -// Forces one route partition to process a second 128-row band. This models a -// maximally skewed router while retaining a noncompact physical route stride. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_partition_stride_case { - %token_count = check.literal value(129) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(1) : index - %route_stride = check.literal value(3) : index - %expert_count = check.literal value(2) : index - %output_size = check.literal value(1) : index - %input = check.generate.fill value(0.00390625) : tensor<129x512xf32> - %q8_input = check.generate.fill value(0) : tensor<129x576xi8> - %route_ids = check.generate.fill value(0) : tensor<129x3xi32> - %expert_table = check.generate.fill value(-1) : tensor<260xi32> - %gate_weight = check.generate.fill value(85) : tensor<2x1x2x144xi8> - %up_weight = check.generate.fill value(-86) : tensor<2x1x2x144xi8> - %expected = check.generate.fill value(0.0) : tensor<129x1x1xf32> - %actual = check.generate.fill value(1.0) : tensor<129x1x1xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<129x512xf32>, tensor<129x576xi8>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<129x3xi32>, tensor<260xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<129x576xi8>, tensor<129x3xi32>, tensor<2x1x2x144xi8>, tensor<2x1x2x144xi8>, tensor<129x1x1xf32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %actual) : [index, index, index, index](index, index, index, index, tensor<129x576xi8>, tensor<260xi32>, tensor<2x1x2x144xi8>, tensor<2x1x2x144xi8>, tensor<129x1x1xf32>) - check.expect.close actual(%actual) expected(%expected) atol(0.25) rtol(9.9999999999999995e-07) nan(same) : tensor<129x1x1xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<[%token_count]x128xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_benchmark_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<8x768xf32> - %completion_counters = check.generate.fill value(0) : tensor<48xi32> - %next_q8_output = check.generate.fill value(1) : tensor<8x864xi8> - %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> - %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> - %expected_counters = check.generate.fill value(0) : tensor<48xi32> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output, %completion_counters, %next_q8_output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>, tensor<48xi32>, tensor<8x864xi8>) - check.expect.close actual(%output) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<8x768xf32> - check.expect.equal actual(%next_q8_output) expected(%expected_q8) : tensor<8x864xi8> - check.expect.equal actual(%completion_counters) expected(%expected_counters) : tensor<48xi32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_benchmark_case { - %token_count = check.literal value(1) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %routed_row_count = check.literal value(8) : index - %q8_input = check.generate.fill value(0) : tensor<2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<8xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<8x768xf32> - %next_q8_output = check.generate.fill value(1) : tensor<8x864xi8> - %expected_output = check.generate.fill value(0.0) : tensor<8x768xf32> - %expected_q8 = check.generate.fill value(0) : tensor<8x864xi8> - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<2304xi8>, tensor<8xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<8x768xf32>) - kernel.launch @ggml_quantize_q8_1_x4_f32[%routed_row_count, %output_size](%routed_row_count, %output_size, %output, %next_q8_output) : [index, index](index, index, tensor<8x768xf32>, tensor<8x864xi8>) - check.expect.close actual(%output) expected(%expected_output) atol(0.0) rtol(0.0) nan(same) : tensor<8x768xf32> - check.expect.equal actual(%next_q8_output) expected(%expected_q8) : tensor<8x864xi8> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - // A period of 127 rotates each physical 128-entry row by one expert while - // keeping the first eight logical routes distinct within every token. - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - // One packed transient buffer holds 128 counts followed by room for 512 - // assignments per expert. The production allocation uses the exact token - // count; this fixed test capacity permits one parameterized benchmark case. - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -// Maximally diverse decode-batch control. Compact route storage makes the -// flattened 0..127 iota assign every M=16 route to a different expert. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %q8_input = check.generate.fill value(0) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 32, 128, 512]) name("token_count") : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %output) : [index, index, index, index, index](index, index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<[%token_count]x128xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512]) name("token_count") : index - %input_size = check.literal value(2048) : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(128) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %q8_input = check.generate.fill value(1) : tensor<[%token_count]x2304xi8> - %route_ids = check.generate.iota offset(0) step(1) period(127) : tensor<[%token_count]x128xi32> - %expert_table = check.generate.fill value(-1) : tensor<65664xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<[%token_count]x2048xf32>, tensor<[%token_count]x2304xi8>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x128xi32>, tensor<65664xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped[%token_count, %route_count, %expert_count, %output_size](%token_count, %route_count, %expert_count, %output_size, %q8_input, %expert_table, %gate_weight, %up_weight, %output) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x2304xi8>, tensor<65664xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_nonzero_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_decode - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8_composed_decode - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_diverse_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_pipeline_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_q8_grouped_pipeline_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom deleted file mode 100644 index 18b915eceef5..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/routed_linear_q4k_f16_wmma.loom +++ /dev/null @@ -1,949 +0,0 @@ -// Gfx11 routed-expert projection matching llama.cpp Vulkan's matrix path. -// -// Each two-wave workgroup computes 64 output channels for 32 routed rows of one -// expert. The waves share padded FP16 operand stages, each owns 32 output -// channels and four WMMA accumulators, and scatters one transposed accumulator -// fragment at a time through a compact route map. Gate and up invoke this same -// projection independently before the separate SwiGLU epilogue. -// -// The raw weights remain [expert][output channel][K / 256][144 bytes]. -// No persistent repacking or expanded-weight allocation is required. -amdgpu.target @qwen3_moe_gfx11_wave64 {subgroup_size = 64} - -amdgpu.target @qwen3_moe_gfx11_wave32 {subgroup_size = 32} - -config.decl @qwen3_moe.routed_gate_up.input_size : %value: index where [range(%value, 512, 32768), mul(%value, 512)] - -// Top-k is a model hyperparameter and is specialized with the kernel. Keeping -// it out of the dynamic workload lets address arithmetic and assignment decode -// fold to the exact model contract. -config.decl @qwen3_moe.routed_gate_up.route_count : %value: index where [range(%value, 1, 8)] - -config.decl @qwen3_moe.routed_gate_up.expert_count : %value: index where [range(%value, 1, 512)] - -config.decl @qwen3_moe.routed_gate_up.output_size : %value: index where [range(%value, 1, 4096)] - -kernel.decl @ggml_quantize_q8_1_x4_f32(%token_count: index, %input_size: index) launch(%token_count: index, %input_size: index, %input: buffer, %output: buffer) - -kernel.decl @qwen3_moe_build_expert_table(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %route_ids: buffer, %expert_table: buffer) - -kernel.decl @qwen3_moe_build_expert_partition_table(%token_count: index, %route_count: index, %expert_count: index) launch(%token_count: index, %route_count: index, %expert_count: index, %expert_table: buffer, %partition_table: buffer) - -func.decl @qwen3_moe_unpack_expert_partition_descriptor(%descriptor: i32) -> (index, index, index) - -kernel.decl @qwen3_moe_routed_gate_up_swiglu_q4k_q8(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index) launch(%token_count: index, %route_count: index, %route_stride: index, %expert_count: index, %output_size: index, %q8_input: buffer, %route_ids: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) - -func.decl @qwen3_moe_q4k_scale_from_header(%scale0: i32, %scale1: i32, %scale2: i32, %q4_group: index) -> (i32, i32) - -// Acquires the packed code word shared by one adjacent Q4_K group pair. -func.def inline @qwen3_moe_q4k_wmma_load_code(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group_pair: index, %packet: index) -> (vector<1xi32>) { - %c8 = index.constant 8 : index - %block_bytes = index.constant 144 : offset - %code_offset = index.constant 16 : offset - %bounded_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index - %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %code_byte_base = index.add %block_byte_base, %code_offset : offset - %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<32xi32> - %q_page = index.mul %bounded_group_pair, %c8 : index - %q_word_index0 = index.add %q_page, %bounded_packet : index - %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index - %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> - func.return %q_word : vector<1xi32> -} - -// Decodes the four adjacent Q4_K values owned by one load packet from an -// already-loaded block header and packed code word. Matrix schedules choose -// the lifetime of both immutable packets. -func.def inline @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group: index, %header_words: vector<4xi32>, %q_word: vector<1xi32>) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %q4_mask = vector.constant 252645135 : vector<1xi32> - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %header_halves = vector.bitcast %header_words : vector<4xi32> to vector<8xf16> - %d_f16 = vector.extract %header_halves[0] : vector<8xf16> -> f16 - %dmin_f16 = vector.extract %header_halves[1] : vector<8xf16> -> f16 - %scale0 = vector.extract %header_words[1] : vector<4xi32> -> i32 - %scale1 = vector.extract %header_words[2] : vector<4xi32> -> i32 - %scale2 = vector.extract %header_words[3] : vector<4xi32> -> i32 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %scale, %minimum = func.call @qwen3_moe_q4k_scale_from_header(%scale0, %scale1, %scale2, %bounded_group) : (i32, i32, i32, index) -> (i32, i32) - %scale_f32 = scalar.uitofp %scale : i32 to f32 - %minimum_f32 = scalar.uitofp %minimum : i32 to f32 - %d_scale = scalar.mulf %d, %scale_f32 : f32 - %minimum_scale = scalar.mulf %dmin, %minimum_f32 : f32 - %q_half = index.rem %bounded_group, %c2 : index - %q_shift_index = index.mul %q_half, %c4 : index - %q_shift_i32 = index.cast %q_shift_index : index to i32 - %q_shift = vector.splat %q_shift_i32 : vector<1xi32> - %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> - %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> - %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> - %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> - // Form adjacent FP16 lanes from fused FP32 affine expressions. AMDGPU maps - // this natural shape to packed mixlo/mixhi instructions where available. - %negative_minimum_scale = scalar.negf %minimum_scale : f32 - %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 - %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 - %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 - %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 - %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 - %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 - %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 - %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 - %half0 = scalar.fptrunc %value0 : f32 to f16 - %half1 = scalar.fptrunc %value1 : f32 to f16 - %half2 = scalar.fptrunc %value2 : f32 to f16 - %half3 = scalar.fptrunc %value3 : f32 to f16 - %result = vector.from_elements %half0, %half1, %half2, %half3 : vector<4xf16> - func.return %result : vector<4xf16> -} - -// Decodes one group when its caller has retained only the Q4_K block header. -func.def inline @qwen3_moe_q4k_wmma_vector4_from_header(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index, %header_words: vector<4xi32>) -> (vector<4xf16>) { - %c2 = index.constant 2 : index - %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index - %q4_group_pair = index.div %bounded_group, %c2 : index - %q_word = func.call @qwen3_moe_q4k_wmma_load_code(%weight, %row_byte_base, %q4_block, %q4_group_pair, %packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - %values = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%bounded_group, %header_words, %q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - func.return %values : vector<4xf16> -} - -// Acquires one naturally aligned Q4_K block header as a single 16-byte packet. -func.def inline @qwen3_moe_q4k_wmma_load_header(%weight: buffer, %row_byte_base: offset, %q4_block: index) -> (vector<4xi32>) { - %c0 = index.constant 0 : index - %block_bytes = index.constant 144 : offset - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %header_view = buffer.view %weight[%block_byte_base] : buffer -> view<4xi32> - %header_words = vector.load %header_view[%c0] : view<4xi32> -> vector<4xi32> - func.return %header_words : vector<4xi32> -} - -// Acquires one block header before decoding the selected four-value group -// packet. Matrix schedules that span several groups call the two operations -// separately so the header lifetime matches their complete block loop. -func.def inline @qwen3_moe_q4k_wmma_vector4(%weight: buffer, %row_byte_base: offset, %q4_block: index, %q4_group: index, %packet: index) -> (vector<4xf16>) { - %header_words = func.call @qwen3_moe_q4k_wmma_load_header(%weight, %row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) - %values = func.call @qwen3_moe_q4k_wmma_vector4_from_header(%weight, %row_byte_base, %q4_block, %q4_group, %packet, %header_words) : (buffer, offset, index, index, index, vector<4xi32>) -> (vector<4xf16>) - func.return %values : vector<4xf16> -} - -kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_linear_q4k_f16_wmma(%token_count: index) { - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %c1 = index.constant 1 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c128 = index.constant 128 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_tiles = index.div %padded_token_count, %c64 : index - kernel.launch.config workgroups(%output_tiles, %route_tiles, %expert_count) workgroup_size(%c128, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %route_tile = kernel.workgroup.id : index - %expert = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 1)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c16 = index.constant 16 : index - %c32 = index.constant 32 : index - %c40 = index.constant 40 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %route_stage_bytes = index.constant 128 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 1024 : offset - %c0_i32 = scalar.constant 0 : i32 - %cn1_i32 = scalar.constant -1 : i32 - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %zero_accumulator = vector.constant 0.0 : vector<8xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %assignment_count = index.mul %token_count, %bounded_route_count : index - %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset - %q4_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.scale %q4_block_count, %q4_block_bytes : index, offset -> offset - %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset - %output_row_count = index.mul %token_count, %bounded_route_count : index - %input_noalias, %expert_table_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %weight, %output : buffer, buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> - %count_view = buffer.view %expert_table_noalias[%c0_offset] : buffer -> view<[%bounded_expert_count]xi32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf32> - %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer - %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer - %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer - %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> - %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> - %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding - %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> - %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> - %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset - %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding - %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> - %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> - %channel_tile_base = index.mul %channel_tile, %c64 : index - %initial_route_tile_base = index.mul %route_tile, %c32 : index - %padded_token_count = index.add %token_count, %c63 : index - %route_partition_count = index.div %padded_token_count, %c64 : index - %route_partition_step = index.mul %route_partition_count, %c32 : index - %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index - %is_workitem_zero = index.cmp eq, %workitem, %c0 : index - %lane_expert_route_count = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %count_view[%bounded_expert] : view<[%bounded_expert_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %expert_route_count_reduced = kernel.workgroup.reduce %lane_expert_route_count : i32 - %expert_route_count_i32 = kernel.subgroup.broadcast.first %expert_route_count_reduced : i32 - %expert_route_count0 = index.cast %expert_route_count_i32 : i32 to index - %expert_route_count = index.assume %expert_route_count0 [range(%expert_route_count0, 0, 2048)] : index - // Route partitions are distributed across the launch grid and continue in - // uniform strides for concentrated routing. Balanced Qwen prefill gives each - // expert 32, 64, or 128 rows at 512, 1024, or 2048 tokens. Concentrated - // experts can consume the full token count without changing the launch - // geometry. - scf.for %route_tile_base = [%initial_route_tile_base to %expert_route_count step %route_partition_step] { - // The first wave snapshots the compact route map once. Every K tile then - // reuses these 32 entries while the full workgroup cooperatively fills the - // operand stages. - %loads_route = index.cmp ult, %workitem, %c32 : index - scf.if %loads_route { - %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index - %assignment_ordinal = index.add %route_tile_base, %local_route : index - %valid_row = index.cmp ult, %assignment_ordinal, %expert_route_count : index - %assignment_i32 = scf.if %valid_row -> (i32) { - %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index - %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %cn1_i32 : i32 - } - view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %load_packet = index.rem %workitem, %c8 : index - %load_k = index.mul %load_packet, %c4 : index - %load_row0 = index.div %workitem, %c8 : index - %load_row = index.assume %load_row0 [range(%load_row0, 0, 15)] : index - %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset - %subgroup_channel_add = index.mul %subgroup, %c32 : index - %subgroup_channel1 = index.add %subgroup_channel_add, %c16 : index - %init00 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init01 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init10 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %init11 = vector.fragment %zero_accumulator shape [%m, %n] : vector<8xf16> - %result00, %result01, %result10, %result11 = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - %block_result00, %block_result01, %block_result10, %block_result11 = scf.for %q4_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 %q4_block, %c256 : index - %group_k_add = index.mul %q4_group, %c32 : index - %k_origin = index.add %block_k_base, %group_k_add : index - scf.for %row_offset = [%c0 to %c64 step %c16] unroll { - %local_row0 = index.add %load_row, %row_offset : index - %local_row = index.assume %local_row0 [range(%local_row0, 0, 63)] : index - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %weight_values = scf.if %valid_channel -> (vector<4xf16>) { - %channel_byte_add = index.scale %channel, %weight_row_bytes : index, offset -> offset - %row_byte_base = index.add %expert_byte_base, %channel_byte_add : offset - %decoded = func.call @qwen3_moe_q4k_wmma_vector4(%weight_noalias, %row_byte_base, %q4_block, %q4_group, %load_packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - %is_activation_row = index.cmp ult, %local_row, %c32 : index - vector.store %weight_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - scf.if %is_activation_row { - %activation_row = index.assume %local_row [range(%local_row, 0, 31)] : index - %assignment_i32 = view.load %route_stage_view[%activation_row] : view<32xi32> -> i32 - %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 - %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %token0 = index.div %bounded_assignment, %bounded_route_count : index - %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> - scf.yield %converted : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %activation_values, %activation_stage_physical_view[%activation_row, %load_k] : vector<4xf16>, view<32x40xf16> - } - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %next00, %next01, %next10, %next11 = scf.for %k_half = [%c0 to %c32 step %c16](%half_acc00 = %acc00 : vector<8xf16>, %half_acc01 = %acc01 : vector<8xf16>, %half_acc10 = %acc10 : vector<8xf16>, %half_acc11 = %acc11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) unroll { - %lhs0 = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %lhs1 = vector.fragment.load %weight_stage_view[%subgroup_channel1, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs0 = vector.fragment.load %activation_fragment_view[%k_half, %c0] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %rhs1 = vector.fragment.load %activation_fragment_view[%k_half, %c16] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %half_next00 = vector.mma %lhs0, %rhs0, %half_acc00 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next01 = vector.mma %lhs0, %rhs1, %half_acc01 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next10 = vector.mma %lhs1, %rhs0, %half_acc10 : vector<16xf16>, vector<16xf16>, vector<8xf16> - %half_next11 = vector.mma %lhs1, %rhs1, %half_acc11 : vector<16xf16>, vector<16xf16>, vector<8xf16> - scf.yield %half_next00, %half_next01, %half_next10, %half_next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %next00, %next01, %next10, %next11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - scf.yield %block_result00, %block_result01, %block_result10, %block_result11 : vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16> - } - // WMMA produces [channel][route] fragments. Store each through a - // transposed, wave-private LDS slice so every lane can scatter four - // contiguous channels to one routed output row. This uses the full wave - // for conversion and publication instead of serializing 16 channels onto - // each of 16 lanes. Subgroup-scoped fences suffice because waves never - // access each other's result slice. - %publish_route0 = index.div %lane, %c4 : index - %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index - %publish_packet0 = index.rem %lane, %c4 : index - %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 3)] : index - %publish_channel_add = index.mul %publish_packet, %c4 : index - %local_route1 = index.add %c16, %publish_route : index - %assignment0_i32 = view.load %route_stage_view[%publish_route] : view<32xi32> -> i32 - %assignment1_i32 = view.load %route_stage_view[%local_route1] : view<32xi32> -> i32 - %assignment0_nonnegative = scalar.cmpi sge, %assignment0_i32, %c0_i32 : i32 - %assignment1_nonnegative = scalar.cmpi sge, %assignment1_i32, %c0_i32 : i32 - %safe_assignment0_i32 = scf.if %assignment0_nonnegative -> (i32) { - scf.yield %assignment0_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %safe_assignment1_i32 = scf.if %assignment1_nonnegative -> (i32) { - scf.yield %assignment1_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %safe_assignment0_0 = index.cast %safe_assignment0_i32 : i32 to index - %safe_assignment1_0 = index.cast %safe_assignment1_i32 : i32 to index - %safe_assignment0 = index.assume %safe_assignment0_0 [range(%safe_assignment0_0, 0, 16383)] : index - %safe_assignment1 = index.assume %safe_assignment1_0 [range(%safe_assignment1_0, 0, 16383)] : index - %bounded_assignment0, %bounded_assignment_count0 = index.assume %safe_assignment0, %assignment_count [lt(%safe_assignment0, %assignment_count)] : index, index - %bounded_assignment1, %bounded_assignment_count1 = index.assume %safe_assignment1, %assignment_count [lt(%safe_assignment1, %assignment_count)] : index, index - %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index - %channel0 = index.add %subgroup_channel_base, %publish_channel_add : index - %channel1_base = index.add %subgroup_channel_base, %c16 : index - %channel1 = index.add %channel1_base, %publish_channel_add : index - %valid_channel0 = index.cmp ult, %channel0, %bounded_output_size : index - %valid_channel1 = index.cmp ult, %channel1, %bounded_output_size : index - %writes00 = scalar.andi %assignment0_nonnegative, %valid_channel0 : i1 - %writes01 = scalar.andi %assignment1_nonnegative, %valid_channel0 : i1 - %writes10 = scalar.andi %assignment0_nonnegative, %valid_channel1 : i1 - %writes11 = scalar.andi %assignment1_nonnegative, %valid_channel1 : i1 - vector.fragment.store %result00, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes00 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment0, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result01, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes01 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel0 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment1, %channel0], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result10, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes10 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment0, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %result11, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<8xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes11 { - %values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<4xf16> - %wide = vector.extf %values : vector<4xf16> to vector<4xf32> - %mask = vector.mask.range [%channel1 to %bounded_output_size step %c1] : index -> vector<4xi1> - vector.store.mask %wide, %output_view[%bounded_assignment1, %channel1], %mask : vector<4xf32>, view<[%output_row_count]x[%bounded_output_size]xf32>, vector<4xi1> - } - // Route and result stages are reused by the next concentrated-routing - // partition. All waves must finish publication before either stage changes. - kernel.barrier scope(workgroup) ordering(acq_rel) - } - kernel.return -} - -// Fuses gate and up projection through SwiGLU for one routed expert tile. -// -// Eight waves split the 64 output channels and 32 routed rows into 16x16 result -// tiles. Each wave carries one gate and one up fragment while the workgroup -// shares one routed activation tile across both contractions. The final -// fragments meet in a wave-private LDS slice and publish the activated product -// directly, avoiding both full-size projection intermediates. -// -// The SwiGLU product is rounded to FP16 at its sole publication point because -// the grouped-down contraction consumes that exact FP16 WMMA operand. Keeping -// the transient in FP16 avoids a widen-store-load-truncate round trip and -// halves its global-memory footprint. -kernel.def target(@qwen3_moe_gfx11_wave32) @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma(%token_count: index) { - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %c1 = index.constant 1 : index - %c31 = index.constant 31 : index - %c32 = index.constant 32 : index - %c63 = index.constant 63 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %padded_output_size = index.add %output_size, %c63 : index - %output_tiles = index.div %padded_output_size, %c64 : index - // The exact partition count lives in device memory. The rounding bound keeps - // it below assignment_partitions + expert_count, so launching the larger - // term lets every workgroup consume at most two strided descriptors. - %assignment_count = index.mul %token_count, %route_count : index - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %expert_count : index - %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %expert_count : index - kernel.launch.config workgroups(%output_tiles, %launch_partition_count, %c1) workgroup_size(%c256, %c1, %c1) : index -} launch(%token_count: index, %input: buffer, %expert_table: buffer, %partition_table: buffer, %gate_weight: buffer, %up_weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %input_size = config.get @qwen3_moe.routed_gate_up.input_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %expert_count = config.get @qwen3_moe.routed_gate_up.expert_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_expert_count = index.assume %expert_count [range(%expert_count, 1, 128)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %channel_tile = kernel.workgroup.id : index - %partition_ordinal = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %subgroup0 = kernel.subgroup.id : index - %subgroup = index.assume %subgroup0 [range(%subgroup0, 0, 7)] : index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %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 - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %c4_bytes = index.constant 4 : offset - %q4_block_bytes = index.constant 144 : offset - %weight_stage_bytes = index.constant 5120 : offset - %activation_stage_bytes = index.constant 2560 : offset - %route_stage_bytes = index.constant 128 : offset - %wave_result_stage_bytes = index.constant 512 : offset - %result_stage_bytes = index.constant 4096 : offset - %c0_i32 = scalar.constant 0 : i32 - %cn1_i32 = scalar.constant -1 : i32 - %c0_i32x1 = vector.constant 0 : vector<1xi32> - %c0_i32x4 = vector.constant 0 : vector<4xi32> - %c0_f16x4 = vector.constant 0.0 : vector<4xf16> - %zero_accumulator = vector.constant 0.0 : vector<16xf16> - %m = index.constant 16 : index - %n = index.constant 16 : index - %k = index.constant 16 : index - %assignment_count = index.mul %token_count, %bounded_route_count : index - %rounded_assignment_count = index.add %assignment_count, %c31 : index - %assignment_partition_count = index.div %rounded_assignment_count, %c32 : index - %maximum_partition_count = index.add %assignment_partition_count, %bounded_expert_count : index - %has_more_assignment_partitions = index.cmp ugt, %assignment_partition_count, %bounded_expert_count : index - %launch_partition_count = scf.select %has_more_assignment_partitions, %assignment_partition_count, %bounded_expert_count : index - %assignment_table_byte_base = index.scale %bounded_expert_count, %c4_bytes : index, offset -> offset - %q4_block_count = index.div %input_size, %c256 : index - %weight_row_bytes = index.scale %q4_block_count, %q4_block_bytes : index, offset -> offset - %weight_expert_bytes = index.scale %bounded_output_size, %weight_row_bytes : index, offset -> offset - %output_row_count = index.mul %token_count, %bounded_route_count : index - %input_noalias, %expert_table_noalias, %partition_table_noalias, %gate_weight_noalias, %up_weight_noalias, %output_noalias = buffer.assume.noalias %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output : buffer, buffer, buffer, buffer, buffer, buffer - %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%token_count]x[%input_size]xf32> - %assignment_view = buffer.view %expert_table_noalias[%assignment_table_byte_base] : buffer -> view<[%bounded_expert_count]x[%token_count]xi32> - %partition_count_view = buffer.view %partition_table_noalias[%c0_offset] : buffer -> view<1xi32> - %partition_descriptor_view = buffer.view %partition_table_noalias[%c4_bytes] : buffer -> view<[%maximum_partition_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%output_row_count]x[%bounded_output_size]xf16> - %weight_stage = buffer.alloca align(16) %weight_stage_bytes : buffer - %activation_stage = buffer.alloca align(16) %activation_stage_bytes : buffer - %route_stage = buffer.alloca align(16) %route_stage_bytes : buffer - %result_stage = buffer.alloca align(16) %result_stage_bytes : buffer - %weight_stage_view = buffer.view %weight_stage[%c0_offset] : buffer -> view<64x40xf16> - %activation_stage_physical_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x40xf16> - %activation_fragment_layout = encoding.layout.strided [1, %c40] : encoding - %activation_fragment_view = buffer.view %activation_stage[%c0_offset] : buffer -> view<32x32xf16, %activation_fragment_layout> - %route_stage_view = buffer.view %route_stage[%c0_offset] : buffer -> view<32xi32> - %wave_result_stage_offset = index.scale %subgroup, %wave_result_stage_bytes : index, offset -> offset - %result_fragment_layout = encoding.layout.strided [1, %c16] : encoding - %result_fragment_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16, %result_fragment_layout> - %result_physical_view = buffer.view %result_stage[%wave_result_stage_offset] : buffer -> view<16x16xf16> - %channel_tile_base = index.mul %channel_tile, %c64 : index - %is_workitem_zero = index.cmp eq, %workitem, %c0 : index - %lane_partition_count_i32 = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %partition_count_view[%c0] : view<1xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %partition_count_reduced = kernel.workgroup.reduce %lane_partition_count_i32 : i32 - %partition_count_i32 = kernel.subgroup.broadcast.first %partition_count_reduced : i32 - %partition_count0 = index.cast %partition_count_i32 : i32 to index - %partition_count, %partition_capacity = index.assume %partition_count0, %maximum_partition_count [lt(%partition_count0, %maximum_partition_count)] : index, index - scf.for %active_partition = [%partition_ordinal to %partition_count step %launch_partition_count] { - %descriptor_ordinal, %descriptor_count = index.assume %active_partition, %partition_count [lt(%active_partition, %partition_count)] : index, index - %table_descriptor_ordinal, %table_descriptor_capacity = index.assume %descriptor_ordinal, %maximum_partition_count [lt(%descriptor_ordinal, %maximum_partition_count)] : index, index - %lane_descriptor_i32 = scf.if %is_workitem_zero -> (i32) { - %loaded = view.load %partition_descriptor_view[%table_descriptor_ordinal] : view<[%maximum_partition_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %c0_i32 : i32 - } - %descriptor_reduced_i32 = kernel.workgroup.reduce %lane_descriptor_i32 : i32 - %descriptor_i32 = kernel.subgroup.broadcast.first %descriptor_reduced_i32 : i32 - %expert, %route_tile_base, %partition_row_count = func.call @qwen3_moe_unpack_expert_partition_descriptor(%descriptor_i32) : (i32) -> (index, index, index) - %bounded_expert, %table_expert_count = index.assume %expert, %bounded_expert_count [lt(%expert, %bounded_expert_count)] : index, index - %loads_route = index.cmp ult, %workitem, %c32 : index - scf.if %loads_route { - %local_route = index.assume %workitem [range(%workitem, 0, 31)] : index - %assignment_ordinal = index.add %route_tile_base, %local_route : index - %valid_row = index.cmp ult, %local_route, %partition_row_count : index - %assignment_i32 = scf.if %valid_row -> (i32) { - %bounded_assignment_ordinal, %table_token_count = index.assume %assignment_ordinal, %token_count [lt(%assignment_ordinal, %token_count)] : index, index - %loaded = view.load %assignment_view[%bounded_expert, %bounded_assignment_ordinal] : view<[%bounded_expert_count]x[%token_count]xi32> -> i32 - scf.yield %loaded : i32 - } else { - scf.yield %cn1_i32 : i32 - } - view.store %assignment_i32, %route_stage_view[%local_route] : i32, view<32xi32> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %load_packet = index.rem %workitem, %c8 : index - %load_k = index.mul %load_packet, %c4 : index - %load_row0 = index.div %workitem, %c8 : index - %load_row = index.assume %load_row0 [range(%load_row0, 0, 31)] : index - %expert_byte_base = index.scale %bounded_expert, %weight_expert_bytes : index, offset -> offset - %channel_subgroup = index.rem %subgroup, %c4 : index - %route_subgroup = index.div %subgroup, %c4 : index - %subgroup_channel_add = index.mul %channel_subgroup, %c16 : index - %subgroup_route_add = index.mul %route_subgroup, %c16 : index - %init_gate = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> - %init_up = vector.fragment %zero_accumulator shape [%m, %n] : vector<16xf16> - %gate_result, %up_result = scf.for %q4_block = [%c0 to %q4_block_count step %c1](%block_gate = %init_gate : vector<16xf16>, %block_up = %init_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { - // Two lanes cover the 64 output rows owned by one load packet. Retain - // both projection headers across the complete eight-group block. - %gate_header0, %gate_header1, %up_header0, %up_header1 = scf.for %header_row_offset = [%c0 to %c64 step %c32](%prior_gate_header0 = %c0_i32x4 : vector<4xi32>, %prior_gate_header1 = %c0_i32x4 : vector<4xi32>, %prior_up_header0 = %c0_i32x4 : vector<4xi32>, %prior_up_header1 = %c0_i32x4 : vector<4xi32>) -> (vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32>) unroll { - %header_local_row0 = index.add %load_row, %header_row_offset : index - %header_local_row = index.assume %header_local_row0 [range(%header_local_row0, 0, 63)] : index - %header_channel = index.add %channel_tile_base, %header_local_row : index - %valid_header_channel = index.cmp ult, %header_channel, %bounded_output_size : index - %loaded_gate_header, %loaded_up_header = scf.if %valid_header_channel -> (vector<4xi32>, vector<4xi32>) { - %header_channel_byte_add = index.scale %header_channel, %weight_row_bytes : index, offset -> offset - %header_row_byte_base = index.add %expert_byte_base, %header_channel_byte_add : offset - %gate_header = func.call @qwen3_moe_q4k_wmma_load_header(%gate_weight_noalias, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) - %up_header = func.call @qwen3_moe_q4k_wmma_load_header(%up_weight_noalias, %header_row_byte_base, %q4_block) : (buffer, offset, index) -> (vector<4xi32>) - scf.yield %gate_header, %up_header : vector<4xi32>, vector<4xi32> - } else { - scf.yield %c0_i32x4, %c0_i32x4 : vector<4xi32>, vector<4xi32> - } - %updates_header0 = index.cmp eq, %header_row_offset, %c0 : index - %next_gate_header0 = scf.select %updates_header0, %loaded_gate_header, %prior_gate_header0 : vector<4xi32> - %next_gate_header1 = scf.select %updates_header0, %prior_gate_header1, %loaded_gate_header : vector<4xi32> - %next_up_header0 = scf.select %updates_header0, %loaded_up_header, %prior_up_header0 : vector<4xi32> - %next_up_header1 = scf.select %updates_header0, %prior_up_header1, %loaded_up_header : vector<4xi32> - scf.yield %next_gate_header0, %next_gate_header1, %next_up_header0, %next_up_header1 : vector<4xi32>, vector<4xi32>, vector<4xi32>, vector<4xi32> - } - %next_block_gate, %next_block_up = scf.for %q4_group_pair = [%c0 to %c4 step %c1](%pair_gate_acc = %block_gate : vector<16xf16>, %pair_up_acc = %block_up : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { - // Each code word supplies the low and high nibbles for one adjacent - // group pair. Retain both projections' words for exactly those uses. - %gate_q_word0, %gate_q_word1, %up_q_word0, %up_q_word1 = scf.for %code_row_offset = [%c0 to %c64 step %c32](%prior_gate_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_gate_q_word1 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word0 = %c0_i32x1 : vector<1xi32>, %prior_up_q_word1 = %c0_i32x1 : vector<1xi32>) -> (vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32>) unroll { - %code_local_row0 = index.add %load_row, %code_row_offset : index - %code_local_row = index.assume %code_local_row0 [range(%code_local_row0, 0, 63)] : index - %code_channel = index.add %channel_tile_base, %code_local_row : index - %valid_code_channel = index.cmp ult, %code_channel, %bounded_output_size : index - %loaded_gate_q_word, %loaded_up_q_word = scf.if %valid_code_channel -> (vector<1xi32>, vector<1xi32>) { - %bounded_q4_group_pair = index.assume %q4_group_pair [range(%q4_group_pair, 0, 3)] : index - %code_channel_byte_add = index.scale %code_channel, %weight_row_bytes : index, offset -> offset - %code_row_byte_base = index.add %expert_byte_base, %code_channel_byte_add : offset - %gate_q_word = func.call @qwen3_moe_q4k_wmma_load_code(%gate_weight_noalias, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - %up_q_word = func.call @qwen3_moe_q4k_wmma_load_code(%up_weight_noalias, %code_row_byte_base, %q4_block, %bounded_q4_group_pair, %load_packet) : (buffer, offset, index, index, index) -> (vector<1xi32>) - scf.yield %gate_q_word, %up_q_word : vector<1xi32>, vector<1xi32> - } else { - scf.yield %c0_i32x1, %c0_i32x1 : vector<1xi32>, vector<1xi32> - } - %updates_q_word0 = index.cmp eq, %code_row_offset, %c0 : index - %next_gate_q_word0 = scf.select %updates_q_word0, %loaded_gate_q_word, %prior_gate_q_word0 : vector<1xi32> - %next_gate_q_word1 = scf.select %updates_q_word0, %prior_gate_q_word1, %loaded_gate_q_word : vector<1xi32> - %next_up_q_word0 = scf.select %updates_q_word0, %loaded_up_q_word, %prior_up_q_word0 : vector<1xi32> - %next_up_q_word1 = scf.select %updates_q_word0, %prior_up_q_word1, %loaded_up_q_word : vector<1xi32> - scf.yield %next_gate_q_word0, %next_gate_q_word1, %next_up_q_word0, %next_up_q_word1 : vector<1xi32>, vector<1xi32>, vector<1xi32>, vector<1xi32> - } - %next_pair_gate, %next_pair_up = scf.for %group_within_pair = [%c0 to %c2 step %c1](%group_gate = %pair_gate_acc : vector<16xf16>, %group_up = %pair_up_acc : vector<16xf16>) -> (vector<16xf16>, vector<16xf16>) { - %q4_group_base = index.mul %q4_group_pair, %c2 : index - %q4_group0 = index.add %q4_group_base, %group_within_pair : index - %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index - %block_k_base = index.mul %q4_block, %c256 : index - %group_k_add = index.mul %q4_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 - %selects_first_row = index.cmp eq, %row_offset, %c0 : index - %selected_gate_header = scf.select %selects_first_row, %gate_header0, %gate_header1 : vector<4xi32> - %selected_gate_q_word = scf.select %selects_first_row, %gate_q_word0, %gate_q_word1 : vector<1xi32> - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %gate_values = scf.if %valid_channel -> (vector<4xf16>) { - %decoded = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group, %selected_gate_header, %selected_gate_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %gate_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - } - %assignment_i32 = view.load %route_stage_view[%load_row] : view<32xi32> -> i32 - %valid_assignment = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 - %activation_values = scf.if %valid_assignment -> (vector<4xf16>) { - %assignment0 = index.cast %assignment_i32 : i32 to index - %assignment = index.assume %assignment0 [range(%assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %assignment, %assignment_count [lt(%assignment, %assignment_count)] : index, index - %token0 = index.div %bounded_assignment, %bounded_route_count : index - %token, %input_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %input_k = index.add %k_origin, %load_k : index - %loaded = vector.load %input_view[%token, %input_k] : view<[%token_count]x[%input_size]xf32> -> vector<4xf32> - %converted = vector.fptrunc %loaded : vector<4xf32> to vector<4xf16> - scf.yield %converted : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %activation_values, %activation_stage_physical_view[%load_row, %load_k] : vector<4xf16>, view<32x40xf16> - kernel.barrier scope(workgroup) ordering(acq_rel) - %gate_next = scf.for %k_half = [%c0 to %c32 step %c16](%half_gate = %group_gate : vector<16xf16>) -> (vector<16xf16>) unroll { - %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %next = vector.mma %lhs, %rhs, %half_gate : vector<16xf16>, vector<16xf16>, vector<16xf16> - scf.yield %next : vector<16xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - 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 - %selects_first_row = index.cmp eq, %row_offset, %c0 : index - %selected_up_header = scf.select %selects_first_row, %up_header0, %up_header1 : vector<4xi32> - %selected_up_q_word = scf.select %selects_first_row, %up_q_word0, %up_q_word1 : vector<1xi32> - %channel = index.add %channel_tile_base, %local_row : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %up_values = scf.if %valid_channel -> (vector<4xf16>) { - %decoded = func.call @qwen3_moe_q4k_wmma_vector4_from_header_code(%q4_group, %selected_up_header, %selected_up_q_word) : (index, vector<4xi32>, vector<1xi32>) -> (vector<4xf16>) - scf.yield %decoded : vector<4xf16> - } else { - scf.yield %c0_f16x4 : vector<4xf16> - } - vector.store %up_values, %weight_stage_view[%local_row, %load_k] : vector<4xf16>, view<64x40xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - %up_next = scf.for %k_half = [%c0 to %c32 step %c16](%half_up = %group_up : vector<16xf16>) -> (vector<16xf16>) unroll { - %lhs = vector.fragment.load %weight_stage_view[%subgroup_channel_add, %k_half] shape [%m, %k] : view<64x40xf16> -> vector<16xf16> - %rhs = vector.fragment.load %activation_fragment_view[%k_half, %subgroup_route_add] shape [%k, %n] : view<32x32xf16, %activation_fragment_layout> -> vector<16xf16> - %next = vector.mma %lhs, %rhs, %half_up : vector<16xf16>, vector<16xf16>, vector<16xf16> - scf.yield %next : vector<16xf16> - } - kernel.barrier scope(workgroup) ordering(acq_rel) - scf.yield %gate_next, %up_next : vector<16xf16>, vector<16xf16> - } - scf.yield %next_pair_gate, %next_pair_up : vector<16xf16>, vector<16xf16> - } - scf.yield %next_block_gate, %next_block_up : vector<16xf16>, vector<16xf16> - } - %publish_route0 = index.div %lane, %c2 : index - %publish_route = index.assume %publish_route0 [range(%publish_route0, 0, 15)] : index - %publish_packet0 = index.rem %lane, %c2 : index - %publish_packet = index.assume %publish_packet0 [range(%publish_packet0, 0, 1)] : index - %publish_channel_add = index.mul %publish_packet, %c8 : index - %local_route = index.add %subgroup_route_add, %publish_route : index - %assignment_i32 = view.load %route_stage_view[%local_route] : view<32xi32> -> i32 - %assignment_nonnegative = scalar.cmpi sge, %assignment_i32, %c0_i32 : i32 - %safe_assignment_i32 = scf.if %assignment_nonnegative -> (i32) { - scf.yield %assignment_i32 : i32 - } else { - scf.yield %c0_i32 : i32 - } - %safe_assignment0 = index.cast %safe_assignment_i32 : i32 to index - %safe_assignment = index.assume %safe_assignment0 [range(%safe_assignment0, 0, 16383)] : index - %bounded_assignment, %bounded_assignment_count = index.assume %safe_assignment, %assignment_count [lt(%safe_assignment, %assignment_count)] : index, index - %subgroup_channel_base = index.add %channel_tile_base, %subgroup_channel_add : index - %channel = index.add %subgroup_channel_base, %publish_channel_add : index - %valid_channel = index.cmp ult, %channel, %bounded_output_size : index - %writes = scalar.andi %assignment_nonnegative, %valid_channel : i1 - vector.fragment.store %gate_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - %gate_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> - %gate_wide = vector.extf %gate_values : vector<8xf16> to vector<8xf32> - %activated = vector.siluf %gate_wide : vector<8xf32> - kernel.barrier scope(subgroup) ordering(acq_rel) - vector.fragment.store %up_result, %result_fragment_view[%c0, %c0] shape [%m, %n] : vector<16xf16>, view<16x16xf16, %result_fragment_layout> - kernel.barrier scope(subgroup) ordering(acq_rel) - scf.if %writes { - %up_values = vector.load %result_physical_view[%publish_route, %publish_channel_add] : view<16x16xf16> -> vector<8xf16> - %up_wide = vector.extf %up_values : vector<8xf16> to vector<8xf32> - %wide_values = vector.mulf %activated, %up_wide : vector<8xf32> - %values = vector.fptrunc %wide_values : vector<8xf32> to vector<8xf16> - %mask = vector.mask.range [%channel to %bounded_output_size step %c1] : index -> vector<8xi1> - vector.store.mask %values, %output_view[%bounded_assignment, %channel], %mask : vector<8xf16>, view<[%output_row_count]x[%bounded_output_size]xf16>, vector<8xi1> - } - kernel.barrier scope(subgroup) ordering(acq_rel) - } - kernel.return -} - -// Applies the model's SwiGLU epilogue to gate and up projections already -// scattered into logical [token][route][output channel] order. -kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_swiglu_f32(%token_count: index) { - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %rounding = index.constant 255 : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %row_count = index.mul %token_count, %route_count : index - %element_count = index.mul %row_count, %output_size : index - %rounded_count = index.add %element_count, %rounding : index - %workgroup_count = index.div %rounded_count, %workgroup_size : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %gate: buffer, %up: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %workgroup = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %workgroup_size = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %row_count = index.mul %token_count, %bounded_route_count : index - %element_count = index.mul %row_count, %bounded_output_size : index - %element = index.madd %workgroup, %workgroup_size, %lane : index - %in_bounds = index.cmp ult, %element, %element_count : index - %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %gate, %up, %output : buffer, buffer, buffer - scf.if %in_bounds { - %bounded_element, %view_element_count = index.assume %element, %element_count [lt(%element, %element_count)] : index, index - %gate_view = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %up_view = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %gate_value = view.load %gate_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %up_value = view.load %up_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %activated = scalar.siluf %gate_value : f32 - %result = scalar.mulf %activated, %up_value : f32 - view.store %result, %output_view[%bounded_element] : f32, view<[%view_element_count]xf32> - } - kernel.return -} - -// Reference FP16 publication path for checking the fused gate/up provider at -// its actual model boundary. Gate and up projections remain independently -// materialized so this schedule does not share the fused kernel's matrix or -// publication implementation. -kernel.def target(@qwen3_moe_gfx11_wave64) @qwen3_moe_routed_swiglu_f16(%token_count: index) { - %c1 = index.constant 1 : index - %workgroup_size = index.constant 256 : index - %rounding = index.constant 255 : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %row_count = index.mul %token_count, %route_count : index - %element_count = index.mul %row_count, %output_size : index - %rounded_count = index.add %element_count, %rounding : index - %workgroup_count = index.div %rounded_count, %workgroup_size : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %gate: buffer, %up: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %route_count = config.get @qwen3_moe.routed_gate_up.route_count : index - %output_size = config.get @qwen3_moe.routed_gate_up.output_size : index - %bounded_route_count = index.assume %route_count [range(%route_count, 1, 8)] : index - %bounded_output_size = index.assume %output_size [range(%output_size, 1, 4096)] : index - %workgroup = kernel.workgroup.id : index - %lane = kernel.workitem.id : index - %workgroup_size = index.constant 256 : index - %c0_offset = index.constant 0 : offset - %row_count = index.mul %token_count, %bounded_route_count : index - %element_count = index.mul %row_count, %bounded_output_size : index - %element = index.madd %workgroup, %workgroup_size, %lane : index - %in_bounds = index.cmp ult, %element, %element_count : index - %gate_noalias, %up_noalias, %output_noalias = buffer.assume.noalias %gate, %up, %output : buffer, buffer, buffer - scf.if %in_bounds { - %bounded_element, %view_element_count = index.assume %element, %element_count [lt(%element, %element_count)] : index, index - %gate_view = buffer.view %gate_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %up_view = buffer.view %up_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%view_element_count]xf16> - %gate_value = view.load %gate_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %up_value = view.load %up_view[%bounded_element] : view<[%view_element_count]xf32> -> f32 - %activated = scalar.siluf %gate_value : f32 - %wide_result = scalar.mulf %activated, %up_value : f32 - %result = scalar.fptrunc %wide_result : f32 to f16 - view.store %result, %output_view[%bounded_element] : f16, view<[%view_element_count]xf16> - } - kernel.return -} - -// The constant input is exactly representable in Q8_1 and FP16. Comparing -// against the established integer-dot path therefore isolates Q4_K decode, -// routing, FP16 matrix accumulation, tail handling, and the SwiGLU boundary. -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential_case { - %token_count = check.literal value(67) : index - %input_size = check.literal value(512) : index - %route_count = check.literal value(2) : index - %route_stride = check.literal value(4) : index - %expert_count = check.literal value(4) : index - %output_size = check.literal value(33) : index - %input = check.generate.fill value(0.00390625) : tensor<67x512xf32> - %q8_input = check.generate.fill value(0) : tensor<67x576xi8> - %route_ids = check.generate.iota offset(0) step(1) period(4) : tensor<67x4xi32> - %expert_table = check.generate.fill value(-1) : tensor<272xi32> - %partition_table = check.generate.fill value(-1) : tensor<10xi32> - %gate_weight = check.generate.fill value(34) : tensor<4x33x2x144xi8> - %up_weight = check.generate.fill value(35) : tensor<4x33x2x144xi8> - %gate_projection = check.generate.fill value(0.0) : tensor<67x2x33xf32> - %up_projection = check.generate.fill value(0.0) : tensor<67x2x33xf32> - %expected = check.generate.fill value(0.0) : tensor<67x2x33xf32> - %actual = check.generate.fill value(1.0) : tensor<67x2x33xf32> - %actual_f16 = check.generate.fill value(1.0) : tensor<67x2x33xf16> - %fused_actual = check.generate.fill value(1.0) : tensor<67x2x33xf16> - kernel.launch @ggml_quantize_q8_1_x4_f32[%token_count, %input_size](%token_count, %input_size, %input, %q8_input) : [index, index](index, index, tensor<67x512xf32>, tensor<67x576xi8>) - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<67x4xi32>, tensor<272xi32>) - kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<272xi32>, tensor<10xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_q8[%token_count, %route_count, %route_stride, %expert_count, %output_size](%token_count, %route_count, %route_stride, %expert_count, %output_size, %q8_input, %route_ids, %gate_weight, %up_weight, %expected) : [index, index, index, index, index](index, index, index, index, index, tensor<67x576xi8>, tensor<67x4xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %gate_weight, %gate_projection) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %up_weight, %up_projection) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<4x33x2x144xi8>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_swiglu_f32[%token_count](%token_count, %gate_projection, %up_projection, %actual) : [index](index, tensor<67x2x33xf32>, tensor<67x2x33xf32>, tensor<67x2x33xf32>) - kernel.launch @qwen3_moe_routed_swiglu_f16[%token_count](%token_count, %gate_projection, %up_projection, %actual_f16) : [index](index, tensor<67x2x33xf32>, tensor<67x2x33xf32>, tensor<67x2x33xf16>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %fused_actual) : [index](index, tensor<67x512xf32>, tensor<272xi32>, tensor<10xi32>, tensor<4x33x2x144xi8>, tensor<4x33x2x144xi8>, tensor<67x2x33xf16>) - check.expect.close actual(%actual) expected(%expected) atol(0.01) rtol(0.01) nan(same) : tensor<67x2x33xf32> - check.expect.close actual(%fused_actual) expected(%actual_f16) atol(0.01) rtol(0.01) nan(same) : tensor<67x2x33xf16> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - // The compact iota exactly matches the oracle's balanced ring: - // expert(token, route) = (8 * token + route) % 128. - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - // One count per expert followed by 2048 assignment slots per expert. - %expert_table = check.generate.fill value(-1) : tensor<262272xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %gate_projection = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %up_projection = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf32> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<262272xi32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %gate_weight, %gate_projection) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - kernel.launch @qwen3_moe_routed_linear_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %up_weight, %up_projection) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf32>) - kernel.launch @qwen3_moe_routed_swiglu_f32[%token_count](%token_count, %gate_projection, %up_projection, %output) : [index](index, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x768xf32>, tensor<[%token_count]x8x768xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf32> - check.return -} - -check.case public @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case { - %token_count = check.param.choice values([1, 2, 4, 8, 16, 17, 32, 63, 128, 129, 512, 1024, 2048]) name("token_count") : index - %route_count = check.literal value(8) : index - %route_stride = check.literal value(8) : index - %expert_count = check.literal value(128) : index - %output_size = check.literal value(768) : index - %input = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - %route_ids = check.generate.iota offset(0) step(1) period(128) : tensor<[%token_count]x8xi32> - // One count per expert followed by 2048 assignment slots per expert. - %expert_table = check.generate.fill value(-1) : tensor<262272xi32> - // One exact count followed by at most 640 packed descriptors. - %partition_table = check.generate.fill value(-1) : tensor<641xi32> - %gate_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %up_weight = check.generate.fill value(0) : tensor<128x768x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x8x768xf16> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x8x768xf16> - kernel.launch @qwen3_moe_build_expert_table[%token_count, %route_count, %route_stride, %expert_count](%token_count, %route_count, %route_stride, %expert_count, %route_ids, %expert_table) : [index, index, index, index](index, index, index, index, tensor<[%token_count]x8xi32>, tensor<262272xi32>) - kernel.launch @qwen3_moe_build_expert_partition_table[%token_count, %route_count, %expert_count](%token_count, %route_count, %expert_count, %expert_table, %partition_table) : [index, index, index](index, index, index, tensor<262272xi32>, tensor<641xi32>) - kernel.launch @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma[%token_count](%token_count, %input, %expert_table, %partition_table, %gate_weight, %up_weight, %output) : [index](index, tensor<[%token_count]x2048xf32>, tensor<262272xi32>, tensor<641xi32>, tensor<128x768x8x144xi8>, tensor<128x768x8x144xi8>, tensor<[%token_count]x8x768xf16>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x8x768xf16> - check.return -} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_differential - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_1024 {token_count = 1024} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_prefill_2048 {token_count = 2048} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_decode {token_count = 1} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_2 {token_count = 2} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_4 {token_count = 4} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_8 {token_count = 8} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_small_batch_16 {token_count = 16} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_17 {token_count = 17} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_32 {token_count = 32} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_63 {token_count = 63} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_128 {token_count = 128} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_129 {token_count = 129} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_512 {token_count = 512} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_1024 {token_count = 1024} - -check.benchmark<@qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_benchmark_case> @qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma_fused_prefill_2048 {token_count = 2048} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom index 699e61ea7fd1..44e464e29634 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_f32.loom @@ -10,29 +10,18 @@ // gfx1151 uses descriptor-backed MUBUF accesses for the decode schedule. The // generic gfx11 provider retains global addressing where descriptor setup costs // more than it saves. +template.decl @qwen3_moe.router_projection.storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) + amdgpu.target @qwen3_moe_router_projection_gfx11_wave32 {subgroup_size = 32} amdgpu.target @qwen3_moe_router_projection_gfx11_wave64 {subgroup_size = 64} -amdgpu.target @qwen3_moe_router_projection_gfx1151 +target.decl @qwen3_moe_router_projection_gfx1151 config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] -// Selects the target's preferred storage contract without exposing addressing -// policy at the projection call site. -func.template target(@qwen3_moe_router_projection_gfx1151) priority(20) @qwen3_moe_router_projection_descriptor_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { - %input_descriptor = buffer.assume.memory_space %input : buffer - %weight_descriptor = buffer.assume.memory_space %weight : buffer - %output_descriptor = buffer.assume.memory_space %output : buffer - func.return %input_descriptor, %weight_descriptor, %output_descriptor : buffer, buffer, buffer -} - -func.template priority(1) @qwen3_moe_router_projection_global_storage(%input: buffer, %weight: buffer, %output: buffer) -> (buffer, buffer, buffer) { - func.return %input, %weight, %output : buffer, buffer, buffer -} - // Scalar differential oracle. It is reachable only from check cases and keeps // production validation independent of the packetized wave schedule. kernel.def target(@qwen3_moe_router_projection_gfx11_wave32) @qwen3_moe_router_projection_f32_reference(%token_count: index) { @@ -95,7 +84,7 @@ kernel.def target(@qwen3_moe_router_projection_gfx11_wave64) @qwen3_moe_router_p %token, %launch_token_count = index.assume %safe_token0, %token_count [lt(%safe_token0, %token_count)] : index, index %expert = index.assume %expert0 [lt(%expert0, %expert_count)] : index %lane_channel = index.mul %lane, %c4 : index - %input_storage, %weight_storage, %output_storage = func.apply(%input, %weight, %output) : (buffer, buffer, buffer) -> (buffer, buffer, buffer) + %input_storage, %weight_storage, %output_storage = template.apply<@qwen3_moe.router_projection.storage>(%input, %weight, %output) : (buffer, buffer, buffer) -> (buffer, buffer, buffer) %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input_storage, %weight_storage, %output_storage : buffer, buffer, buffer %input_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%hidden_size]xf32> %weight_view = buffer.view %weight_noalias[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom index b3c68f8f817e..b542f9219452 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_projection_top8_fused_f32.loom @@ -8,106 +8,33 @@ // // This schedule is decode-only. Prefill has enough token parallelism to keep // projection and routing independent without paying completion atomics. -amdgpu.target @qwen3_moe_router_fused_gfx11_wave64 {subgroup_size = 64} +template.decl @qwen3_moe.router.fused.experts_per_wave() -> (index) -amdgpu.target @qwen3_moe_router_fused_gfx1151 +template.decl @qwen3_moe.router.fused.projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) -config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] +template.decl @qwen3_moe.router.top8.row(%arg0: i1, %arg1: index, %arg2: index, %arg3: index, %arg4: buffer, %arg5: buffer, %arg6: buffer) -config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] +template.decl @qwen3_moe.router_projection.storage(%arg0: buffer, %arg1: buffer, %arg2: buffer) -> (buffer, buffer, buffer) -kernel.decl @qwen3_moe_router_projection_f32_four_row_wave32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +amdgpu.target @qwen3_moe_router_fused_gfx11_wave64 {subgroup_size = 64} -kernel.decl @qwen3_moe_router_top8_f32(%token_count: index, %route_id_stride: index) launch(%token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) +target.decl @qwen3_moe_router_fused_gfx1151 -func.template target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_experts_per_wave() -> (index) { - %c2 = index.constant 2 : index - func.return %c2 : index -} +config.decl @qwen3_moe.model.hidden_size : %value: index where [range(%value, 128, 32768), mul(%value, 128)] -func.template requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_experts_per_wave() -> (index) { - %c4 = index.constant 4 : index - func.return %c4 : index -} +config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] -func.template device target(@qwen3_moe_router_fused_gfx1151) requires [#target.subgroup.size<64>] priority(20) @qwen3_moe_router_fused_two_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { - %c0_offset = index.constant 0 : offset - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %expert_base = index.mul %expert_tile, %c2 : index - %expert1 = index.add %expert_base, %c1 : index - %lane_channel = index.mul %lane, %c4 : index - %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> - %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> - %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> - %c0_f32x2 = vector.splat %c0_f32 : vector<2xf32> - %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x2 : vector<2xf32>) -> (vector<2xf32>) unroll(%c4) schedule(interleaved) { - %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %accumulator0 = vector.extract %accumulators[0] : vector<2xf32> -> f32 - %accumulator1 = vector.extract %accumulators[1] : vector<2xf32> -> f32 - %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 - %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 - %next_accumulators = vector.from_elements %next0, %next1 : vector<2xf32> - scf.yield %next_accumulators : vector<2xf32> - } - %sums = kernel.subgroup.reduce %lane_sums : vector<2xf32> - scf.if %writes_projection { - vector.store %sums, %logits_view[%token, %expert_base] : vector<2xf32>, view<[%token_count]x[%expert_count]xf32> - } - func.return -} +kernel.decl @qwen3_moe_router_projection_f32_four_row_wave32(%token_count$26: index) launch(%token_count$27: index, %input: buffer, %weight: buffer, %output: buffer) -func.template device requires [#target.subgroup.size<64>] priority(1) @qwen3_moe_router_fused_four_expert_projection(%hidden_size: index, %expert_count: index, %token_count: index, %token: index, %expert_tile: index, %lane: index, %writes_projection: i1, %input: buffer, %weight: buffer, %logits: buffer) { - %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 - %c256 = index.constant 256 : index - %c0_f32 = scalar.constant 0.0 : f32 - %expert_base = index.mul %expert_tile, %c4 : index - %expert1 = index.add %expert_base, %c1 : index - %expert2 = index.add %expert_base, %c2 : index - %expert3 = index.add %expert_base, %c3 : index - %lane_channel = index.mul %lane, %c4 : index - %input_view = buffer.view %input[%c0_offset] : buffer -> view<[%token_count]x[%hidden_size]xf32> - %weight_view = buffer.view %weight[%c0_offset] : buffer -> view<[%expert_count]x[%hidden_size]xf32> - %logits_view = buffer.view %logits[%c0_offset] : buffer -> view<[%token_count]x[%expert_count]xf32> - %c0_f32x4 = vector.splat %c0_f32 : vector<4xf32> - %lane_sums = scf.for %channel = [%lane_channel to %hidden_size step %c256](%accumulators = %c0_f32x4 : vector<4xf32>) -> (vector<4xf32>) unroll(%c4) schedule(interleaved) { - %input_values = vector.load %input_view[%token, %channel] : view<[%token_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values0 = vector.load %weight_view[%expert_base, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values1 = vector.load %weight_view[%expert1, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values2 = vector.load %weight_view[%expert2, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %weight_values3 = vector.load %weight_view[%expert3, %channel] : view<[%expert_count]x[%hidden_size]xf32> -> vector<4xf32> - %accumulator0 = vector.extract %accumulators[0] : vector<4xf32> -> f32 - %accumulator1 = vector.extract %accumulators[1] : vector<4xf32> -> f32 - %accumulator2 = vector.extract %accumulators[2] : vector<4xf32> -> f32 - %accumulator3 = vector.extract %accumulators[3] : vector<4xf32> -> f32 - %next0 = vector.dotf %input_values, %weight_values0, %accumulator0 : vector<4xf32>, vector<4xf32>, f32 - %next1 = vector.dotf %input_values, %weight_values1, %accumulator1 : vector<4xf32>, vector<4xf32>, f32 - %next2 = vector.dotf %input_values, %weight_values2, %accumulator2 : vector<4xf32>, vector<4xf32>, f32 - %next3 = vector.dotf %input_values, %weight_values3, %accumulator3 : vector<4xf32>, vector<4xf32>, f32 - %next_accumulators = vector.from_elements %next0, %next1, %next2, %next3 : vector<4xf32> - scf.yield %next_accumulators : vector<4xf32> - } - %sums = kernel.subgroup.reduce %lane_sums : vector<4xf32> - scf.if %writes_projection { - vector.store %sums, %logits_view[%token, %expert_base] : vector<4xf32>, view<[%token_count]x[%expert_count]xf32> - } - func.return -} +kernel.decl @qwen3_moe_router_top8_f32(%token_count$31: index, %route_id_stride$32: index) launch(%token_count$33: index, %route_id_stride$34: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) + +kernel.decl @qwen3_moe_router_top8_f32_check_one() launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) kernel.def target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projection_top8_fused_decode_f32(%token_count: index, %route_id_stride: index) { %expert_count = config.get @qwen3_moe.router.expert_count : index %c1 = index.constant 1 : index %wave_size = target.subgroup.size : index - %experts_per_wave = func.apply() : () -> (index) + %experts_per_wave = index.constant 4 : index %expert_tiles = index.div %expert_count, %experts_per_wave : index kernel.launch.config workgroups(%expert_tiles, %c1, %c1) workgroup_size(%wave_size, %c1, %c1) : index } launch(%token_count: index, %route_id_stride: index, %input: buffer, %weight: buffer, %logits: buffer, %completion_counter: buffer, %route_ids: buffer, %route_weights: buffer) { @@ -129,7 +56,7 @@ kernel.def target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projec %c0_offset = index.constant 0 : offset %counter_scratch_bytes = index.constant 4 : offset %token = index.assume %c0 [lt(%c0, %bounded_token_count)] : index - %input_storage, %weight_storage, %logits_storage = func.apply(%input, %weight, %logits) : (buffer, buffer, buffer) -> (buffer, buffer, buffer) + %input_storage, %weight_storage, %logits_storage = template.apply<@qwen3_moe.router_projection.storage>(%input, %weight, %logits) : (buffer, buffer, buffer) -> (buffer, buffer, buffer) %input_noalias, %weight_noalias, %logits_noalias, %completion_counter_noalias, %route_ids_noalias, %route_weights_noalias = buffer.assume.noalias %input_storage, %weight_storage, %logits_storage, %completion_counter, %route_ids, %route_weights : buffer, buffer, buffer, buffer, buffer, buffer %completion_counter_aligned = buffer.assume.alignment %completion_counter_noalias {minimum_alignment = 16} : buffer %completion_counter_view = buffer.view %completion_counter_aligned[%c0_offset] : buffer -> view<1xi32> @@ -138,7 +65,7 @@ kernel.def target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projec %lane_i32 = index.cast %lane : index to i32 %writes_projection = scalar.cmpi eq, %lane_i32, %c0_i32 : i32 - func.apply(%hidden_size, %expert_count, %bounded_token_count, %token, %expert_tile, %lane, %writes_projection, %input_noalias, %weight_noalias, %logits_noalias) : (index, index, index, index, index, index, i1, buffer, buffer, buffer) + template.apply<@qwen3_moe.router.fused.projection>(%hidden_size, %expert_count, %bounded_token_count, %token, %expert_tile, %lane, %writes_projection, %input_noalias, %weight_noalias, %logits_noalias) : (index, index, index, index, index, index, i1, buffer, buffer, buffer) // Every projection store precedes its workgroup's release. The last arrival // acquires all preceding releases before any lane loads the complete row. @@ -154,7 +81,7 @@ kernel.def target(@qwen3_moe_router_fused_gfx11_wave64) @qwen3_moe_router_projec %negative_expert_tile_count_i32 = scalar.subi %c0_i32, %expert_tile_count_i32 : i32 %is_last_projection = scalar.cmpi eq, %old_counter, %last_expert_tile_i32 : i32 scf.if %is_last_projection { - func.apply(%publishes_row, %token, %bounded_token_count, %bounded_route_id_stride, %logits_noalias, %route_ids_noalias, %route_weights_noalias) : (i1, index, index, index, buffer, buffer, buffer) + template.apply<@qwen3_moe.router.top8.row>(%publishes_row, %token, %bounded_token_count, %bounded_route_id_stride, %logits_noalias, %route_ids_noalias, %route_weights_noalias) : (i1, index, index, index, buffer, buffer, buffer) // The counter cannot become reusable until all route stores complete. kernel.barrier scope(workgroup) ordering(acq_rel) scf.if %writes_projection { @@ -183,7 +110,7 @@ check.case public @qwen3_moe_router_projection_top8_fused_differential_case { %completion_counter = check.generate.fill value(0) : tensor<1xi32> %expected_counter = check.generate.fill value(0) : tensor<1xi32> kernel.launch @qwen3_moe_router_projection_f32_four_row_wave32[%token_count](%token_count, %input, %weight, %expected_logits) : [index](index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>) - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %expected_logits, %expected_route_ids, %expected_route_weights) : [index, index](index, index, tensor<1x128xf32>, tensor<1x8xi32>, tensor<1x8xf32>) + kernel.launch @qwen3_moe_router_top8_f32_check_one(%expected_logits, %expected_route_ids, %expected_route_weights) : (tensor<1x128xf32>, tensor<1x8xi32>, tensor<1x8xf32>) kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %actual_logits0, %completion_counter, %actual_route_ids0, %actual_route_weights0) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) kernel.launch @qwen3_moe_router_projection_top8_fused_decode_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %input, %weight, %actual_logits1, %completion_counter, %actual_route_ids1, %actual_route_weights1) : [index, index](index, index, tensor<1x2048xf32>, tensor<128x2048xf32>, tensor<1x128xf32>, tensor<1xi32>, tensor<1x8xi32>, tensor<1x8xf32>) check.expect.close actual(%actual_logits0) expected(%expected_logits) atol(0.001) rtol(0.001) nan(same) : tensor<1x128xf32> @@ -222,6 +149,8 @@ check.case public @qwen3_moe_router_projection_top8_composed_benchmark_case { check.return } +check.benchmark<@qwen3_moe_router_projection_top8_fused_differential_case> @qwen3_moe_router_projection_top8_fused_differential + check.benchmark<@qwen3_moe_router_projection_top8_fused_benchmark_case> @qwen3_moe_router_projection_top8_fused_decode check.benchmark<@qwen3_moe_router_projection_top8_composed_benchmark_case> @qwen3_moe_router_projection_top8_composed_decode diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom index e2b2908d72c2..88e404695f0c 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_moe/qwen3_moe/router_top8_f32.loom @@ -16,104 +16,14 @@ // Selection therefore operates directly on logits and only the selected // values reach the exponential. This preserves the model contract while // avoiding exponentials for experts that cannot contribute to the result. +template.decl @qwen3_moe.router.top8.row(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) + amdgpu.target @qwen3_moe_router_gfx11_wave64 {subgroup_size = 64} config.decl @qwen3_moe.router.expert_count : %value: index where [range(%value, 32, 512), mul(%value, 32)] config.decl @qwen3_moe.router.route_count : %value: index where [range(%value, 1, 32)] -// Selects and normalizes one logical router row. The caller owns mapping a -// subgroup to a safe physical row and tells the helper whether that row should -// publish. Keeping row semantics independent of launch geometry lets fused -// producers consume this exact tie-breaking and normalization contract. -func.template device requires [#target.subgroup.size<64>] @qwen3_moe_router_top8_row_f32(%valid_token: i1, %token: index, %token_count: index, %route_id_stride: index, %logits: buffer, %route_ids: buffer, %route_weights: buffer) { - %expert_count0 = config.get @qwen3_moe.router.expert_count : index - %route_count0 = config.get @qwen3_moe.router.route_count : index - %bounded_token_count = index.assume %token_count [range(%token_count, 1, 2048)] : index - %bounded_route_id_stride0 = index.assume %route_id_stride [range(%route_id_stride, 1, 512)] : index - %expert_count, %route_count, %bounded_route_id_stride = index.assume %expert_count0, %route_count0, %bounded_route_id_stride0 [range(%expert_count0, 32, 512), mul(%expert_count0, 32), range(%route_count0, 1, 32), le(%route_count0, %expert_count0), le(%route_count0, %bounded_route_id_stride0)] : index, index, index - %safe_token, %launch_token_count = index.assume %token, %bounded_token_count [lt(%token, %bounded_token_count)] : index, index - %lane = kernel.subgroup.lane.id : index - %c0 = index.constant 0 : index - %c1 = index.constant 1 : index - %wave_size = target.subgroup.size : index - %negative_large = scalar.constant -3.4028234663852885e+38 : f32 - %largest_i32 = scalar.constant 2147483647 : i32 - %c0_offset = index.constant 0 : offset - %experts_per_lane0 = index.div %expert_count, %wave_size : index - %experts_per_lane = index.assume %experts_per_lane0 [range(%experts_per_lane0, 1, 16)] : index - %lane_expert_base = index.mul %lane, %experts_per_lane : index - %route_id_storage_count = index.mul %launch_token_count, %bounded_route_id_stride : index - %route_weight_storage_count = index.mul %launch_token_count, %route_count : index - %logits_noalias, %route_ids_noalias, %route_weights_noalias = buffer.assume.noalias %logits, %route_ids, %route_weights : buffer, buffer, buffer - %logits_view = buffer.view %logits_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x[%expert_count]xf32> - %route_ids_view = buffer.view %route_ids_noalias[%c0_offset] : buffer -> view<[%route_id_storage_count]xi32> - %route_weights_view = buffer.view %route_weights_noalias[%c0_offset] : buffer -> view<[%route_weight_storage_count]xf32> - %initial_logits = vector.load %logits_view[%safe_token, %lane_expert_base] : view<[%launch_token_count]x[%expert_count]xf32> -> vector<[%experts_per_lane]xf32> - %remaining_final, %selected_logit = scf.for %route = [%c0 to %route_count step %c1](%remaining_logits = %initial_logits : vector<[%experts_per_lane]xf32>, %lane_selected_logit = %negative_large : f32) -> (vector<[%experts_per_lane]xf32>, f32) { - %local_value, %local_id = scf.for %slot = [%c0 to %experts_per_lane step %c1](%best_value = %negative_large : f32, %best_id = %largest_i32 : i32) -> (f32, i32) unroll { - %candidate_value = vector.extract %remaining_logits[%slot] : vector<[%experts_per_lane]xf32> -> f32 - %candidate_expert = index.add %lane_expert_base, %slot : index - %candidate_id = index.cast %candidate_expert : index to i32 - %is_greater = scalar.cmpf ogt, %candidate_value, %best_value : f32 - %is_equal = scalar.cmpf oeq, %candidate_value, %best_value : f32 - %is_lower_id = scalar.cmpi ult, %candidate_id, %best_id : i32 - %is_lower_tie = scalar.andi %is_equal, %is_lower_id : i1 - %is_better = scalar.ori %is_greater, %is_lower_tie : i1 - %next_value, %next_id = scf.if %is_better -> (f32, i32) { - scf.yield %candidate_value, %candidate_id : f32, i32 - } else { - scf.yield %best_value, %best_id : f32, i32 - } - scf.yield %next_value, %next_id : f32, i32 - } - %winner_value = kernel.subgroup.reduce %local_value : f32 - %matches_winner = scalar.cmpf oeq, %local_value, %winner_value : f32 - %winner_lane_mask = kernel.subgroup.vote.ballot %matches_winner : i1 -> i64 - %nonzero_winner_lane_mask = scalar.assume %winner_lane_mask [ne(%winner_lane_mask, 0)] : i64 - %winner_lane_i64 = scalar.cttzi %nonzero_winner_lane_mask : i64 - %winner_lane0 = index.cast %winner_lane_i64 : i64 to index - %winner_lane = index.assume %winner_lane0 [range(%winner_lane0, 0, 63)] : index - %local_expert0 = index.cast %local_id : i32 to index - %local_expert = index.assume %local_expert0 [range(%local_expert0, 0, 511)] : index - %winner_slot = index.rem %local_expert, %experts_per_lane : index - %owns_winner = index.cmp eq, %lane, %winner_lane : index - %next_remaining_logits = scf.if %owns_winner -> (vector<[%experts_per_lane]xf32>) { - %removed = vector.insert %negative_large into %remaining_logits[%winner_slot] : f32, vector<[%experts_per_lane]xf32> - scf.yield %removed : vector<[%experts_per_lane]xf32> - } else { - scf.yield %remaining_logits : vector<[%experts_per_lane]xf32> - } - %publishes_route_id = scalar.andi %valid_token, %owns_winner : i1 - scf.if %publishes_route_id { - %route_id_token_base = index.mul %safe_token, %bounded_route_id_stride : index - %route_id_index = index.add %route_id_token_base, %route : index - view.store %local_id, %route_ids_view[%route_id_index] : i32, view<[%route_id_storage_count]xi32> - } - %lane_publishes_route = index.cmp eq, %lane, %route : index - %publishes_route = scalar.andi %valid_token, %lane_publishes_route : i1 - %next_lane_selected_logit = scf.if %publishes_route -> (f32) { - scf.yield %winner_value : f32 - } else { - scf.yield %lane_selected_logit : f32 - } - scf.yield %next_remaining_logits, %next_lane_selected_logit : vector<[%experts_per_lane]xf32>, f32 - } - %selected_max = kernel.subgroup.reduce %selected_logit : f32 - %selected_delta = scalar.subf %selected_logit, %selected_max : f32 - %unnormalized_weight = scalar.expf %selected_delta : f32 - %selected_sum = kernel.subgroup.reduce %unnormalized_weight : f32 - %route_weight = scalar.divf %unnormalized_weight, %selected_sum : f32 - %lane_publishes_weight = index.cmp ult, %lane, %route_count : index - %publishes_weight = scalar.andi %valid_token, %lane_publishes_weight : i1 - scf.if %publishes_weight { - %route_weight_token_base = index.mul %safe_token, %route_count : index - %route_weight_index = index.add %route_weight_token_base, %lane : index - view.store %route_weight, %route_weights_view[%route_weight_index] : f32, view<[%route_weight_storage_count]xf32> - } - func.return -} - kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32(%token_count: index, %route_id_stride: index) { %c1 = index.constant 1 : index %c3 = index.constant 3 : index @@ -132,7 +42,7 @@ kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32(%to %token0 = index.add %token_base, %subgroup : index %valid_token = index.cmp ult, %token0, %token_count : index %safe_token = scf.select %valid_token, %token0, %c0 : index - func.apply(%valid_token, %safe_token, %token_count, %route_id_stride, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %token_count, %route_id_stride, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) kernel.return } @@ -160,6 +70,54 @@ kernel.def @qwen3_moe_router_top8_wide_stride_reference() { kernel.return } +// Fixed-shape wrappers keep the correctness cases independent of dynamic +// subrange proof without changing the production kernel's lowering. +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32_check_one() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c1 = index.constant 1 : index + %c8 = index.constant 8 : index + %valid_token = index.cmp ult, %subgroup, %c1 : index + %safe_token = scf.select %valid_token, %subgroup, %c0 : index + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %c1, %c8, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32_check_two_compact() { + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c8 = index.constant 8 : index + %valid_token = index.cmp ult, %subgroup, %c2 : index + %safe_token = scf.select %valid_token, %subgroup, %c0 : index + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %c2, %c8, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + +kernel.def target(@qwen3_moe_router_gfx11_wave64) @qwen3_moe_router_top8_f32_check_two_wide() { + %c1 = index.constant 1 : index + %c256 = index.constant 256 : index + kernel.launch.config workgroups(%c1, %c1, %c1) workgroup_size(%c256, %c1, %c1) : index +} launch(%logits: buffer, %route_ids: buffer, %route_weights: buffer) { + %subgroup = kernel.subgroup.id : index + %c0 = index.constant 0 : index + %c2 = index.constant 2 : index + %c16 = index.constant 16 : index + %valid_token = index.cmp ult, %subgroup, %c2 : index + %safe_token = scf.select %valid_token, %subgroup, %c0 : index + template.apply<@qwen3_moe.router.top8.row>(%valid_token, %safe_token, %c2, %c16, %logits, %route_ids, %route_weights) : (i1, index, index, index, buffer, buffer, buffer) + kernel.return +} + // Period-eight logits place sixteen equal maxima in different lane-local // slots. Repeated selection must remove each winner and retain the lower-ID // half of that tie set. Two rows leave two inactive waves in the four-wave @@ -177,9 +135,9 @@ check.case public @qwen3_moe_router_top8_f32_repeated_maxima_case { %expected_ids = check.generate.iota offset(7) step(8) period(8) : tensor<2x8xi32> %expected_wide_ids = check.generate.fill value(-1) : tensor<2x16xi32> %expected_weights = check.generate.fill value(0.125) : tensor<2x8xf32> - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %route_id_stride](%token_count, %route_id_stride, %logits, %route_ids, %route_weights) : [index, index](index, index, tensor<2x128xf32>, tensor<2x8xi32>, tensor<2x8xf32>) + kernel.launch @qwen3_moe_router_top8_f32_check_two_compact(%logits, %route_ids, %route_weights) : (tensor<2x128xf32>, tensor<2x8xi32>, tensor<2x8xf32>) kernel.launch @qwen3_moe_router_top8_wide_stride_reference(%expected_wide_ids) : (tensor<2x16xi32>) - kernel.launch @qwen3_moe_router_top8_f32[%token_count, %wide_route_id_stride](%token_count, %wide_route_id_stride, %logits, %wide_route_ids, %wide_route_weights) : [index, index](index, index, tensor<2x128xf32>, tensor<2x16xi32>, tensor<2x8xf32>) + kernel.launch @qwen3_moe_router_top8_f32_check_two_wide(%logits, %wide_route_ids, %wide_route_weights) : (tensor<2x128xf32>, tensor<2x16xi32>, tensor<2x8xf32>) check.expect.equal actual(%route_ids) expected(%expected_ids) : tensor<2x8xi32> check.expect.close actual(%route_weights) expected(%expected_weights) atol(1.0000000000000001e-05) rtol(1.0000000000000001e-05) nan(same) : tensor<2x8xf32> check.expect.equal actual(%wide_route_ids) expected(%expected_wide_ids) : tensor<2x16xi32> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_q4k.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_q4k.loom deleted file mode 100644 index 1bce11c86515..000000000000 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/qwen_owned/token_embedding_q4k.loom +++ /dev/null @@ -1,191 +0,0 @@ -// Copyright 2026 The IREE Authors -// -// Licensed under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -// Gathers token rows directly from the model's unmodified GGUF Q4_K payload -// and decodes them to the owned F32 hidden-state layout. The fixed model row -// contract is: -// Q4_K: [vocabulary row][8 blocks][144 bytes] -> [2048xf32] -amdgpu.target @qwen_token_embedding_gfx11_wave64 {subgroup_size = 64} - -kernel.def target(@qwen_token_embedding_gfx11_wave64) export("qwen_token_embedding_q4k") @qwen_token_embedding_q4k(%token_count: index, %vocabulary_count: index) { - %c1 = index.constant 1 : index - %c2 = index.constant 2 : index - %workgroup_count = index.mul %token_count, %c2 : index - %workgroup_size = index.constant 256 : index - kernel.launch.config workgroups(%workgroup_count, %c1, %c1) workgroup_size(%workgroup_size, %c1, %c1) : index -} launch(%token_count: index, %vocabulary_count: index, %token_ids: buffer, %weight: buffer, %output: buffer) where [range(%token_count, 1, 2048)] { - %bounded_vocabulary_count = index.assume %vocabulary_count [range(%vocabulary_count, 1, 262144)] : index - %workgroup = kernel.workgroup.id : index - %workitem = kernel.workitem.id : index - %c0 = index.constant 0 : index - %c2 = index.constant 2 : index - %c4 = index.constant 4 : index - %c8 = index.constant 8 : index - %c64 = index.constant 64 : index - %c256 = index.constant 256 : index - %packets_per_token = index.constant 512 : index - %c0_offset = index.constant 0 : offset - %block_bytes = index.constant 144 : offset - %scale_offset = index.constant 4 : offset - %code_offset = index.constant 16 : offset - %row_bytes = index.constant 1152 : offset - %c2_i32 = scalar.constant 2 : i32 - %c4_i32 = scalar.constant 4 : i32 - %c15 = vector.constant 15 : vector<1xi32> - %c48 = vector.constant 48 : vector<1xi32> - %q4_mask = vector.constant 252645135 : vector<1xi32> - %packet0 = index.madd %workgroup, %c256, %workitem : index - %packet_count = index.mul %token_count, %packets_per_token : index - %packet, %launch_packet_count = index.assume %packet0, %packet_count [lt(%packet0, %packet_count)] : index, index - %token0 = index.div %packet, %packets_per_token : index - %token, %launch_token_count = index.assume %token0, %token_count [lt(%token0, %token_count)] : index, index - %row_packet0 = index.rem %packet, %packets_per_token : index - %row_packet = index.assume %row_packet0 [range(%row_packet0, 0, 511)] : index - %q4_block0 = index.div %row_packet, %c64 : index - %q4_block = index.assume %q4_block0 [range(%q4_block0, 0, 7)] : index - %block_packet0 = index.rem %row_packet, %c64 : index - %block_packet = index.assume %block_packet0 [range(%block_packet0, 0, 63)] : index - %q4_group0 = index.div %block_packet, %c8 : index - %q4_group = index.assume %q4_group0 [range(%q4_group0, 0, 7)] : index - %load_packet0 = index.rem %block_packet, %c8 : index - %load_packet = index.assume %load_packet0 [range(%load_packet0, 0, 7)] : index - %token_ids_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %token_ids, %weight, %output : buffer, buffer, buffer - %token_ids_view = buffer.view %token_ids_noalias[%c0_offset] : buffer -> view<[%launch_token_count]xi32> - %output_view = buffer.view %output_noalias[%c0_offset] : buffer -> view<[%launch_token_count]x2048xf32> - // Token IDs are validated at the request boundary. This trusted consumer - // carries that fact into row addressing instead of adding a fallback row. - %token_id_raw = view.load %token_ids_view[%token] : view<[%launch_token_count]xi32> -> i32 - %token_id_i32 = scalar.assume %token_id_raw [range(%token_id_raw, 0, 262143)] : i32 - %token_id0 = index.cast %token_id_i32 : i32 to index - %token_id, %weight_vocabulary_count = index.assume %token_id0, %bounded_vocabulary_count [lt(%token_id0, %bounded_vocabulary_count)] : index, index - %row_byte_base = index.scale %token_id, %row_bytes : index, offset -> offset - %block_byte_add = index.scale %q4_block, %block_bytes : index, offset -> offset - %block_byte_base = index.add %row_byte_base, %block_byte_add : offset - %scale_byte_base = index.add %block_byte_base, %scale_offset : offset - %code_byte_base = index.add %block_byte_base, %code_offset : offset - %dm_view = buffer.view %weight_noalias[%block_byte_base] : buffer -> view<2xf16> - %scale_view = buffer.view %weight_noalias[%scale_byte_base] : buffer -> view<3xi32> - %code_view = buffer.view %weight_noalias[%code_byte_base] : buffer -> view<32xi32> - %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> - %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> - %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 - %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 - %d = scalar.extf %d_f16 : f16 to f32 - %dmin = scalar.extf %dmin_f16 : f16 to f32 - %q_page0 = index.div %q4_group, %c2 : index - %q_page = index.mul %q_page0, %c8 : index - %q_word_index0 = index.add %q_page, %load_packet : index - %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index - %is_low = index.cmp ult, %q4_group, %c4 : index - %scale_lane = index.rem %q4_group, %c4 : index - %scale_shift_index = index.mul %scale_lane, %c8 : index - %scale_shift_i32 = index.cast %scale_shift_index : index to i32 - %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> - %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 - %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 - %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 - %scale0 = vector.splat %scale0_i32 : vector<1xi32> - %scale1 = vector.splat %scale1_i32 : vector<1xi32> - %scale2 = vector.splat %scale2_i32 : vector<1xi32> - %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 - %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 - %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> - %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> - %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 - %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 - %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> - %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> - %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> - %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> - %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> - %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> - %scale = vector.ori %scale_low, %scale_high : vector<1xi32> - %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> - %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> - %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> - %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> - %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> - %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> - %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> - %d_vector1 = vector.splat %d : vector<1xf32> - %dmin_vector1 = vector.splat %dmin : vector<1xf32> - %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> - %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> - %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 - %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 - %q_word = vector.load %code_view[%q_word_index] : view<32xi32> -> vector<1xi32> - %q_half = index.rem %q4_group, %c2 : index - %q_shift_index = index.mul %q_half, %c4 : index - %q_shift_i32 = index.cast %q_shift_index : index to i32 - %q_shift = vector.splat %q_shift_i32 : vector<1xi32> - %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> - %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> - %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> - %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> - %negative_minimum_scale = scalar.negf %minimum_scale : f32 - %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 - %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 - %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 - %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 - %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 - %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 - %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 - %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 - %values = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> - %output_channel = index.mul %row_packet, %c4 : index - vector.store %values, %output_view[%token, %output_channel] : vector<4xf32>, view<[%launch_token_count]x2048xf32> - kernel.return -} - -// Uniform 0x55 Q4_K bytes encode d=dmin=85.3125, scale=minimum=21, -// and q=5 in every group, producing 85.3125 * 21 * (5 - 1) = 7166.25. -check.case public @qwen_token_embedding_q4k_decode_case { - %token_count = check.literal value(1) : index - %vocabulary_count = check.literal value(1) : index - %token_ids = check.generate.fill value(0) : tensor<1xi32> - %weight = check.generate.fill value(85) : tensor<1x8x144xi8> - %output = check.generate.fill value(0.0) : tensor<1x2048xf32> - %expected = check.generate.fill value(7166.25) : tensor<1x2048xf32> - kernel.launch @qwen_token_embedding_q4k[%token_count, %vocabulary_count](%token_count, %vocabulary_count, %token_ids, %weight, %output) : [index, index](index, index, tensor<1xi32>, tensor<1x8x144xi8>, tensor<1x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<1x2048xf32> - check.return -} - -// Reverse-order IDs exercise the first, interior, and final legal row while -// remaining suitable for access-sanitized execution. -check.case public @qwen_token_embedding_q4k_row_access_case { - %token_count = check.literal value(3) : index - %vocabulary_count = check.literal value(3) : index - %token_ids = check.generate.iota offset(2) step(-1) : tensor<3xi32> - %weight = check.generate.fill value(0) : tensor<3x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<3x2048xf32> - %expected = check.generate.fill value(0.0) : tensor<3x2048xf32> - kernel.launch @qwen_token_embedding_q4k[%token_count, %vocabulary_count](%token_count, %vocabulary_count, %token_ids, %weight, %output) : [index, index](index, index, tensor<3xi32>, tensor<3x8x144xi8>, tensor<3x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<3x2048xf32> - check.return -} - -check.case public @qwen_token_embedding_q4k_benchmark_case { - %token_count = check.param.choice values([32, 128, 512]) name("token_count") : index - %vocabulary_count = check.literal value(151936) : index - %token_ids = check.generate.iota offset(151424) step(1) period(512) : tensor<[%token_count]xi32> - %weight = check.generate.fill value(0) : tensor<151936x8x144xi8> - %output = check.generate.fill value(1.0) : tensor<[%token_count]x2048xf32> - %expected = check.generate.fill value(0.0) : tensor<[%token_count]x2048xf32> - kernel.launch @qwen_token_embedding_q4k[%token_count, %vocabulary_count](%token_count, %vocabulary_count, %token_ids, %weight, %output) : [index, index](index, index, tensor<[%token_count]xi32>, tensor<151936x8x144xi8>, tensor<[%token_count]x2048xf32>) - check.expect.close actual(%output) expected(%expected) atol(0.0) rtol(0.0) nan(same) : tensor<[%token_count]x2048xf32> - check.return -} - -check.benchmark<@qwen_token_embedding_q4k_decode_case> @qwen_token_embedding_q4k_decode - -check.benchmark<@qwen_token_embedding_q4k_row_access_case> @qwen_token_embedding_q4k_row_access - -check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_32 {token_count = 32} - -check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_128 {token_count = 128} - -check.benchmark<@qwen_token_embedding_q4k_benchmark_case> @qwen_token_embedding_q4k_prefill_512 {token_count = 512} diff --git a/ggml/src/ggml-hrx/loom-jit.cpp b/ggml/src/ggml-hrx/loom-jit.cpp index 81e57ca07321..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 = {}; } @@ -67,7 +83,6 @@ struct ggml_hrx_loom_jit_amdgpu { loomc_target_profile_t * target_profile = nullptr; loomc_compiler_t * compiler = nullptr; loomc_pass_program_t * pass_program = nullptr; - loomc_amdgpu_runtime_global_flags_t runtime_globals = LOOMC_AMDGPU_RUNTIME_GLOBAL_NONE; }; namespace { @@ -219,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); } @@ -361,15 +375,40 @@ void * ggml_hrx_loom_jit_malloc_copy(const void * data, size_t size, bool nul_te return result; } -const loomc_artifact_t * ggml_hrx_loom_jit_find_artifact(const loomc_result_t * result, - loomc_artifact_kind_t kind, - loomc_string_view_t format) { +#if defined(LOOMC_ARTIFACT_ROLE_COMPILE_REPORT) +using ggml_hrx_loom_jit_artifact_selector_t = loomc_string_view_t; +#define GGML_HRX_LOOM_ARTIFACT_COMPILE_REPORT loomc_make_cstring_view(LOOMC_ARTIFACT_ROLE_COMPILE_REPORT) +#define GGML_HRX_LOOM_ARTIFACT_MODULE loomc_make_cstring_view(LOOMC_ARTIFACT_ROLE_MODULE) +#define GGML_HRX_LOOM_ARTIFACT_LAUNCH_CONFIG loomc_make_cstring_view(LOOMC_ARTIFACT_ROLE_LAUNCH_CONFIG) +#define GGML_HRX_LOOM_ARTIFACT_KERNEL loomc_make_cstring_view(LOOMC_ARTIFACT_ROLE_KERNEL) +#define GGML_HRX_LOOM_ARTIFACT_MANIFEST loomc_make_cstring_view(LOOMC_ARTIFACT_ROLE_ARTIFACT_MANIFEST) +static bool ggml_hrx_loom_jit_artifact_matches(const loomc_artifact_t * artifact, + ggml_hrx_loom_jit_artifact_selector_t selector) { + return loomc_string_view_equal(artifact->role, selector); +} +#else +using ggml_hrx_loom_jit_artifact_selector_t = loomc_artifact_kind_t; +#define GGML_HRX_LOOM_ARTIFACT_COMPILE_REPORT LOOMC_ARTIFACT_KIND_REPORT +#define GGML_HRX_LOOM_ARTIFACT_MODULE LOOMC_ARTIFACT_KIND_MODULE +#define GGML_HRX_LOOM_ARTIFACT_LAUNCH_CONFIG LOOMC_ARTIFACT_KIND_LAUNCH_CONFIG +#define GGML_HRX_LOOM_ARTIFACT_KERNEL LOOMC_ARTIFACT_KIND_EXECUTABLE +#define GGML_HRX_LOOM_ARTIFACT_MANIFEST LOOMC_ARTIFACT_KIND_REPORT +static bool ggml_hrx_loom_jit_artifact_matches(const loomc_artifact_t * artifact, + ggml_hrx_loom_jit_artifact_selector_t selector) { + return artifact->kind == selector; +} +#endif + +const loomc_artifact_t * ggml_hrx_loom_jit_find_artifact(const loomc_result_t * result, + ggml_hrx_loom_jit_artifact_selector_t selector, + loomc_string_view_t format) { for (loomc_host_size_t i = 0; i < loomc_result_artifact_count(result); ++i) { const loomc_artifact_t * artifact = loomc_result_artifact_at(result, i); if (!artifact) { continue; } - if (artifact->kind == kind && loomc_string_view_equal(artifact->format, format)) { + if (ggml_hrx_loom_jit_artifact_matches(artifact, selector) && + loomc_string_view_equal(artifact->format, format)) { return artifact; } } @@ -389,30 +428,44 @@ hrx_status_t ggml_hrx_loom_jit_copy_artifact_bytes(const loomc_artifact_t * arti if (!artifact || !out_data || !out_size) { return hrx_ok_status(); } - void * copy = ggml_hrx_loom_jit_malloc_copy(artifact->contents.data, artifact->contents.data_length, nul_terminate); + + loomc_byte_span_t contents = loomc_byte_span_empty(); + const bool contents_borrowed = loomc_byte_sequence_try_get_contiguous_span(artifact->contents, &contents); + if (!contents_borrowed) { + loomc_status_t status = loomc_byte_sequence_clone(artifact->contents, loomc_allocator_system(), &contents); + if (!loomc_status_is_ok(status)) { + return ggml_hrx_loom_jit_status_from_loom(status, "copy Loom artifact"); + } + } + + void * copy = ggml_hrx_loom_jit_malloc_copy(contents.data, contents.data_length, nul_terminate); + if (!contents_borrowed) { + loomc_allocator_free(loomc_allocator_system(), const_cast(contents.data)); + } if (!copy) { return ggml_hrx_loom_jit_make_status(HRX_STATUS_OUT_OF_MEMORY, "failed to copy Loom artifact"); } *out_data = copy; - *out_size = artifact->contents.data_length; + *out_size = contents.data_length; 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, nullptr, nullptr, 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"); } @@ -421,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); @@ -438,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"); } @@ -513,18 +583,19 @@ hrx_status_t ggml_hrx_loom_jit_parse_sanitizer_reporting(const char * return ggml_hrx_loom_jit_make_status(HRX_STATUS_INVALID_ARGUMENT, message); } -loomc_amdgpu_runtime_global_flags_t ggml_hrx_loom_jit_runtime_globals(loomc_sanitizer_checks_t sanitizer_checks) { - if (!sanitizer_checks) { - return LOOMC_AMDGPU_RUNTIME_GLOBAL_NONE; - } - loomc_amdgpu_runtime_global_flags_t runtime_globals = LOOMC_AMDGPU_RUNTIME_GLOBAL_FEEDBACK_CONFIG; - if (sanitizer_checks & LOOMC_SANITIZER_CHECK_ACCESS) { - runtime_globals |= LOOMC_AMDGPU_RUNTIME_GLOBAL_ASAN_CONFIG; - } - return runtime_globals; +} // namespace + +void ggml_hrx_loom_jit_launch_program_release(ggml_hrx_loom_jit_launch_program * program) { + delete program; } -} // namespace +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) { @@ -592,7 +663,6 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_create(const ggml_hrx_loom_jit_amdgpu_opti return sanitizer_reporting_status; } } - jit->runtime_globals = ggml_hrx_loom_jit_runtime_globals(sanitizer_options.checks); loomc_target_pipeline_options_t pipeline_options = {}; pipeline_options.type = LOOMC_STRUCTURE_TYPE_TARGET_PIPELINE_OPTIONS; pipeline_options.structure_size = sizeof(pipeline_options); @@ -675,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, @@ -749,15 +819,7 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * dependency_sources.push_back(dependency_source); loomc_link_index_source_options_t dependency_link_options = {}; dependency_link_options.provider_name = loomc_make_cstring_view(dependency.source_identifier); - if (dependency.source_format == GGML_HRX_LOOM_JIT_SOURCE_FORMAT_BYTECODE) { - dependency_link_options.role = LOOMC_LINK_PROVIDER_ROLE_INPUT; - } else { - const std::string dependency_text(static_cast(dependency.source_data), - dependency.source_size); - dependency_link_options.role = dependency_text.find("config.decl") != std::string::npos ? - LOOMC_LINK_PROVIDER_ROLE_INPUT : - LOOMC_LINK_PROVIDER_ROLE_LIBRARY; - } + dependency_link_options.role = LOOMC_LINK_PROVIDER_ROLE_INPUT; status = loomc_link_index_builder_add_source(link_index_builder.get(), dependency_source, &dependency_link_options, nullptr); if (!loomc_status_is_ok(status)) { @@ -796,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 = {}; @@ -826,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); @@ -874,16 +937,29 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * } result.reset(); } + const loomc_target_specialization_t specialization = { + loomc_make_cstring_view(options->root_symbol), + jit->target_profile, + }; + loomc_target_specialization_options_t compile_target_options = {}; + compile_target_options.type = LOOMC_STRUCTURE_TYPE_TARGET_SPECIALIZATION_OPTIONS; + compile_target_options.structure_size = sizeof(compile_target_options); + compile_target_options.specializations = &specialization; + compile_target_options.specialization_count = 1; + const loomc_string_view_t root_symbols[] = { loomc_make_cstring_view(options->root_symbol) }; loomc_link_options_t link_options = {}; link_options.type = LOOMC_STRUCTURE_TYPE_LINK_OPTIONS; link_options.structure_size = sizeof(link_options); - link_options.next = nullptr; + link_options.next = &compile_target_options; + link_options.mode = LOOMC_LINK_MODE_LINK; link_options.link_index = link_index.get(); link_options.module_name = loomc_make_cstring_view(options->module_name); link_options.root_symbols = root_symbols; link_options.root_symbol_count = 1; link_options.flags = LOOMC_LINK_FLAG_STRIP_TEST_SYMBOLS; + link_options.config.bindings = config_bindings.get(); + link_options.config.binding_count = options->config_binding_count; status = loomc_link_module(linker.get(), workspace.get(), &link_options, module.out(), result.out()); if (!loomc_status_is_ok(status)) { return ggml_hrx_loom_jit_status_from_loom(status, "link Loom root"); @@ -893,27 +969,16 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * } result.reset(); - const loomc_target_specialization_t specialization = { - loomc_make_cstring_view(options->root_symbol), - jit->target_profile, - }; - loomc_target_specialization_options_t compile_target_options = {}; - compile_target_options.type = LOOMC_STRUCTURE_TYPE_TARGET_SPECIALIZATION_OPTIONS; - compile_target_options.structure_size = sizeof(compile_target_options); - compile_target_options.specializations = &specialization; - compile_target_options.specialization_count = 1; loomc_compile_options_t compile_options = {}; compile_options.type = LOOMC_STRUCTURE_TYPE_COMPILE_OPTIONS; compile_options.structure_size = sizeof(compile_options); 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.bindings = config_bindings.get(); - compile_options.config.binding_count = options->config_binding_count; - compile_options.config.flags = LOOMC_CONFIG_POLICY_FLAG_REQUIRE_RESOLVED; + compile_options.config_flags = LOOMC_CONFIG_POLICY_FLAG_REQUIRE_RESOLVED; status = loomc_compile_module(jit->compiler, workspace.get(), jit->pass_program, module.get(), &compile_options, loomc_allocator_system(), result.out()); if (!loomc_status_is_ok(status)) { @@ -924,45 +989,49 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * } hrx_status_t hrx_status = hrx_ok_status(); - const loomc_artifact_t * compile_report = ggml_hrx_loom_jit_find_artifact( - result.get(), LOOMC_ARTIFACT_KIND_REPORT, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_JSON)); + const loomc_artifact_t * compile_report = + ggml_hrx_loom_jit_find_artifact(result.get(), GGML_HRX_LOOM_ARTIFACT_COMPILE_REPORT, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_JSON)); hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes(compile_report, reinterpret_cast(&out_result->compile_report_json), &out_result->compile_report_json_size, true); if (!hrx_status_is_ok(hrx_status)) { return hrx_status; } - const loomc_artifact_t * final_module = ggml_hrx_loom_jit_find_artifact( - result.get(), LOOMC_ARTIFACT_KIND_MODULE, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_LOOM_TEXT)); + const loomc_artifact_t * final_module = + ggml_hrx_loom_jit_find_artifact(result.get(), GGML_HRX_LOOM_ARTIFACT_MODULE, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_LOOM_TEXT)); hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes(final_module, reinterpret_cast(&out_result->final_module_text), &out_result->final_module_text_size, true); 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; } const loomc_artifact_t * launch_config = - ggml_hrx_loom_jit_find_artifact(result.get(), LOOMC_ARTIFACT_KIND_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(); - loomc_amdgpu_emit_options_t amdgpu_options = {}; - amdgpu_options.type = LOOMC_STRUCTURE_TYPE_AMDGPU_EMIT_OPTIONS; - amdgpu_options.structure_size = sizeof(amdgpu_options); - amdgpu_options.next = nullptr; - amdgpu_options.runtime_globals = jit->runtime_globals; const loomc_option_entry_t emit_entries[] = { { loomc_make_cstring_view(LOOMC_EMIT_OPTION_KEY_IDENTIFIER), @@ -972,7 +1041,7 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * loomc_option_dict_t option_dict = {}; option_dict.type = LOOMC_STRUCTURE_TYPE_OPTION_DICT; option_dict.structure_size = sizeof(option_dict); - option_dict.next = &amdgpu_options; + option_dict.next = nullptr; option_dict.entries = emit_entries; option_dict.entry_count = options->artifact_identifier ? 1 : 0; loomc_artifact_manifest_options_t manifest_options = {}; @@ -1004,8 +1073,9 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * return hrx_status; } - const loomc_artifact_t * hsaco = ggml_hrx_loom_jit_find_artifact( - result.get(), LOOMC_ARTIFACT_KIND_EXECUTABLE, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_AMDGPU_HSACO)); + const loomc_artifact_t * hsaco = + ggml_hrx_loom_jit_find_artifact(result.get(), GGML_HRX_LOOM_ARTIFACT_KERNEL, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_AMDGPU_HSACO)); if (!hsaco) { out_result->reset(); return ggml_hrx_loom_jit_make_status(HRX_STATUS_NOT_FOUND, "Loom did not return an AMDGPU HSACO artifact"); @@ -1013,16 +1083,16 @@ hrx_status_t ggml_hrx_loom_jit_amdgpu_compile(ggml_hrx_loom_jit_amdgpu * hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes(hsaco, &out_result->hsaco_data, &out_result->hsaco_size, false); if (hrx_status_is_ok(hrx_status)) { const loomc_artifact_t * report = - ggml_hrx_loom_jit_find_artifact(result.get(), LOOMC_ARTIFACT_KIND_REPORT, + ggml_hrx_loom_jit_find_artifact(result.get(), GGML_HRX_LOOM_ARTIFACT_COMPILE_REPORT, loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_COMPILE_REPORT_JSON)); hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes(report, reinterpret_cast(&out_result->compile_report_json), &out_result->compile_report_json_size, true); } if (hrx_status_is_ok(hrx_status)) { - const loomc_artifact_t * manifest = - ggml_hrx_loom_jit_find_artifact(result.get(), LOOMC_ARTIFACT_KIND_REPORT, - loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_ARTIFACT_MANIFEST_JSON)); + const loomc_artifact_t * manifest = ggml_hrx_loom_jit_find_artifact( + result.get(), GGML_HRX_LOOM_ARTIFACT_MANIFEST, + loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_ARTIFACT_MANIFEST_JSON)); hrx_status = ggml_hrx_loom_jit_copy_artifact_bytes( manifest, reinterpret_cast(&out_result->manifest_json), &out_result->manifest_json_size, true); } 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 3aefad8d3297..8c30489611a2 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.cpp @@ -8,12 +8,57 @@ #include "runtime/transient-arena.h" #include +#include +#include +#include +#include +#include +#include +#include +#include #include #include #include namespace ggml::hrx { +GraphReplayStreamState::~GraphReplayStreamState() { + clear(); +} + +void GraphReplayStreamState::clear() { + for (PendingHostWriteback & writeback : pending_host_writebacks_) { + if (writeback.retained_buffer != nullptr) { + hrx_buffer_release(writeback.retained_buffer); + writeback.retained_buffer = nullptr; + } + } + pending_host_writebacks_.clear(); +} + +void GraphReplayStreamState::mark_stream_synchronized() { + for (PendingHostWriteback & writeback : pending_host_writebacks_) { + if (writeback.host_destination != nullptr && writeback.mapped_source != nullptr && writeback.size > 0) { + std::memcpy(writeback.host_destination, writeback.mapped_source, writeback.size); + } + if (writeback.retained_buffer != nullptr) { + hrx_buffer_release(writeback.retained_buffer); + writeback.retained_buffer = nullptr; + } + } + pending_host_writebacks_.clear(); +} + +void GraphReplayStreamState::add_host_writeback(void * host_destination, + const void * mapped_source, + size_t size, + hrx_buffer_t buffer) { + if (buffer != nullptr) { + hrx_buffer_retain(buffer); + } + pending_host_writebacks_.push_back({ host_destination, mapped_source, size, buffer }); +} + PreparedProgramConstantBuffer::~PreparedProgramConstantBuffer() { if (buffer != nullptr) { hrx_buffer_release(buffer); @@ -28,8 +73,8 @@ PreparedProgramConstantBuffer::PreparedProgramConstantBuffer(PreparedProgramCons other.size = 0; } -PreparedProgramConstantBuffer & -PreparedProgramConstantBuffer::operator=(PreparedProgramConstantBuffer && other) noexcept { +PreparedProgramConstantBuffer & PreparedProgramConstantBuffer::operator=( + PreparedProgramConstantBuffer && other) noexcept { if (this != &other) { if (buffer != nullptr) { hrx_buffer_release(buffer); @@ -50,11 +95,19 @@ RecordedCommandGraph::~RecordedCommandGraph() { if (graph != nullptr) { hrx_graph_release(graph); } + for (hrx_buffer_t buffer : retained_buffers) { + hrx_buffer_release(buffer); + } + for (hrx_executable_t executable : retained_executables) { + hrx_executable_release(executable); + } } RecordedCommandGraph::RecordedCommandGraph(RecordedCommandGraph && other) noexcept : graph(std::exchange(other.graph, nullptr)), exec(std::exchange(other.exec, nullptr)), + retained_buffers(std::move(other.retained_buffers)), + retained_executables(std::move(other.retained_executables)), bound_transient_arena_allocation_id(other.bound_transient_arena_allocation_id), dispatch_count(other.dispatch_count), status(std::move(other.status)) { @@ -70,8 +123,16 @@ RecordedCommandGraph & RecordedCommandGraph::operator=(RecordedCommandGraph && o if (graph != nullptr) { hrx_graph_release(graph); } + for (hrx_buffer_t buffer : retained_buffers) { + hrx_buffer_release(buffer); + } + for (hrx_executable_t executable : retained_executables) { + hrx_executable_release(executable); + } graph = std::exchange(other.graph, nullptr); exec = std::exchange(other.exec, nullptr); + retained_buffers = std::move(other.retained_buffers); + retained_executables = std::move(other.retained_executables); bound_transient_arena_allocation_id = other.bound_transient_arena_allocation_id; dispatch_count = other.dispatch_count; status = std::move(other.status); @@ -87,10 +148,361 @@ static const char * status_first_error(const Status & status) { return status.errors().empty() ? "" : status.errors().front().c_str(); } +static bool environment_flag_enabled(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0'); +} + +static const char * environment_string_value(const char * name) { + const char * value = std::getenv(name); + return value != nullptr && value[0] != '\0' && !(value[0] == '0' && value[1] == '\0') ? value : nullptr; +} + +static size_t environment_size_value(const char * name, size_t fallback) { + const char * value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return fallback; + } + return static_cast(std::strtoull(value, nullptr, 0)); +} + +static void debug_serial_log(const char * event, const char * phase, const std::string & detail) { + std::fprintf(stderr, "hrx debug serial: %s phase=%s", event, phase); + if (!detail.empty()) { + std::fprintf(stderr, " %s", detail.c_str()); + } + std::fprintf(stderr, "\n"); + std::fflush(stderr); +} + +class DebugSerialExecutionTrace { + public: + explicit DebugSerialExecutionTrace(const CommandProgramExecutionContext & context) : + context_(context), + enabled_(debug_serial_command_execution_enabled()) {} + + DebugSerialExecutionTrace(const CommandProgramExecutionContext & context, const CommandProgram & commands) : + DebugSerialExecutionTrace(context) { + if (enabled_) { + std::ostringstream out; + out << "main_commands=" << commands.commands.size() + << " init_commands=" << commands.initialization_commands.size() + << " transient_arena=" << commands.transients.arena_size; + program_detail_ = out.str(); + } + } + + DebugSerialExecutionTrace(const CommandProgramExecutionContext & context, + const PreparedCommandProgram & commands) : + DebugSerialExecutionTrace(context) { + if (enabled_) { + std::ostringstream out; + out << "main_commands=" << commands.commands.size() + << " init_commands=" << commands.initialization_commands.size(); + program_detail_ = out.str(); + } + } + + bool enabled() const { + return enabled_; + } + + void log_program_begin() const { + log("begin", "command-program", program_detail_); + } + + void log_program_end(bool success) const { + log(success ? "end" : "failed", "command-program", program_detail_); + } + + bool sync_program_phase(const char * phase) const { + return sync(phase, program_detail_); + } + + bool sync(const char * phase, const std::string & detail = {}) const { + if (!enabled_) { + return true; + } + log("sync-begin", phase, detail); + if (ErrorResult error = take_status(hrx_stream_synchronize(context_.stream))) { + GGML_LOG_ERROR("HRX debug serial sync failed phase=%s %s: %s\n", phase, detail.c_str(), error->c_str()); + return false; + } + log("sync-end", phase, detail); + return true; + } + + bool should_trace_program(size_t command_count) const { + const char * value = std::getenv("GGML_HRX_DEBUG_SERIAL_PROGRAM_COMMANDS"); + return value == nullptr || value[0] == '\0' || command_count == std::strtoull(value, nullptr, 0); + } + + bool should_trace_command(size_t program_command_count, size_t index) const { + if (!should_trace_program(program_command_count)) { + return false; + } + const size_t start = environment_size_value("GGML_HRX_DEBUG_SERIAL_START_COMMAND", 0); + const size_t end = environment_size_value("GGML_HRX_DEBUG_SERIAL_END_COMMAND", + std::numeric_limits::max()); + return index >= start && index <= end; + } + + void log(const char * event, const char * phase, const std::string & detail = {}) const { + if (enabled_) { + debug_serial_log(event, phase, detail); + } + } + + static std::string command_detail(const char * list_kind, + size_t list_size, + size_t index, + const PreparedCommand & command) { + std::ostringstream out; + out << "list=" << list_kind + << " index=" << index + << " list_commands=" << list_size + << " ordinal=" << command.ordinal + << " kind=" << command_kind_name(command.kind); + if (command.kind == CommandKind::Kernel) { + out << " kernel_id=" << command.kernel.specialization.kernel_id + << " bindings=" << command.kernel.bindings.size(); + } + return out.str(); + } + + static std::string list_detail(const char * list_kind, size_t list_size, size_t program_command_count) { + std::ostringstream out; + out << "list=" << list_kind + << " list_commands=" << list_size + << " main_commands=" << program_command_count; + return out.str(); + } + + static std::string host_staging_detail(const HostStagingBuffer & staging) { + std::ostringstream out; + out << "value=" << staging.value << " length=" << staging.length; + return out.str(); + } + + private: + const CommandProgramExecutionContext & context_; + bool enabled_ = false; + std::string program_detail_; +}; + +#define HRX_DEBUG_SERIAL_SYNC(trace, phase) \ + do { \ + if (!(trace).sync_program_phase(phase)) { \ + return false; \ + } \ + } while (0) + static bool resource_access_writes(ResourceAccess access) { return access == ResourceAccess::Write || access == ResourceAccess::ReadWrite; } +static bool resource_access_reads(ResourceAccess access) { + return access == ResourceAccess::Read || access == ResourceAccess::ReadWrite; +} + +static bool dispatch_trace_enabled() { + return environment_string_value("GGML_HRX_TRACE_DISPATCH_DIR") != nullptr; +} + +static bool dispatch_trace_blob_capture_enabled() { + return environment_flag_enabled("GGML_HRX_TRACE_BLOBS"); +} + +static bool dispatch_trace_should_trace_program(size_t program_command_count) { + const char * value = environment_string_value("GGML_HRX_TRACE_PROGRAM_COMMANDS"); + return value == nullptr || program_command_count == std::strtoull(value, nullptr, 0); +} + +static bool dispatch_trace_should_trace_command(size_t program_command_count, size_t command_index) { + if (!dispatch_trace_should_trace_program(program_command_count)) { + return false; + } + const size_t start = environment_size_value("GGML_HRX_TRACE_START_COMMAND", 0); + const size_t end = + environment_size_value("GGML_HRX_TRACE_END_COMMAND", std::numeric_limits::max()); + return command_index >= start && command_index <= end; +} + +static uint64_t dispatch_trace_hash_bytes(const uint8_t * data, size_t size) { + uint64_t hash = UINT64_C(1469598103934665603); + for (size_t i = 0; i < size; ++i) { + hash ^= static_cast(data[i]); + hash *= UINT64_C(1099511628211); + } + return hash; +} + +static std::string dispatch_trace_hex(uint64_t value) { + std::ostringstream out; + out << std::hex << std::setw(16) << std::setfill('0') << value; + return out.str(); +} + +static std::string dispatch_trace_json_escape(const std::string & text) { + std::string escaped; + escaped.reserve(text.size()); + for (const unsigned char c : text) { + switch (c) { + case '\\': escaped += "\\\\"; break; + case '"': escaped += "\\\""; break; + case '\b': escaped += "\\b"; break; + case '\f': escaped += "\\f"; break; + case '\n': escaped += "\\n"; break; + case '\r': escaped += "\\r"; break; + case '\t': escaped += "\\t"; break; + default: + if (c < 0x20) { + static constexpr char kHex[] = "0123456789abcdef"; + escaped += "\\u00"; + escaped.push_back(kHex[c >> 4]); + escaped.push_back(kHex[c & 0xf]); + } else { + escaped.push_back(static_cast(c)); + } + break; + } + } + return escaped; +} + +static std::string dispatch_trace_kernel_name(const CommandProgramExecutionContext & context, + const PreparedCommand & command) { + if (context.corpus != nullptr && context.target != nullptr) { + const KernelResolveResult resolved = + resolve_kernel_definition(*context.corpus, context.target, command.kernel.specialization.kernel_id); + if (resolved.found()) { + return kernel_definition_name(*resolved.definition); + } + } + std::ostringstream out; + out << "kernel_id=" << command.kernel.specialization.kernel_id; + return out.str(); +} + +static Status dispatch_trace_download_binding(const CommandProgramExecutionContext & context, + const PreparedCommandBinding & binding, + std::vector & data) { + Status status; + data.resize(binding.ref.length); + if (binding.ref.length == 0) { + return status; + } + if (context.host_transfers == nullptr) { + status.log("missing HRX host transfer manager for dispatch trace"); + return status; + } + status.append(context.host_transfers->download_synchronous(context.stream, binding.ref.buffer, binding.ref.offset, + data.data(), binding.ref.length)); + return status; +} + +static Status dispatch_trace_write_blob(const std::filesystem::path & path, const std::vector & data) { + Status status; + try { + std::filesystem::create_directories(path.parent_path()); + std::ofstream out(path, std::ios::binary | std::ios::trunc); + if (!out) { + status.log("failed to create dispatch trace blob %s", path.string().c_str()); + return status; + } + out.write(reinterpret_cast(data.data()), static_cast(data.size())); + } catch (const std::exception & error) { + status.log("failed to write dispatch trace blob %s: %s", path.string().c_str(), error.what()); + } + return status; +} + +static Status dispatch_trace_command_bindings(const CommandProgramExecutionContext & context, + const PreparedCommand & command, + const char * list_kind, + size_t list_size, + size_t command_index, + size_t program_command_count, + const char * stage) { + Status status; + const char * directory_value = environment_string_value("GGML_HRX_TRACE_DISPATCH_DIR"); + if (directory_value == nullptr || !dispatch_trace_should_trace_command(program_command_count, command_index)) { + return status; + } + const bool before_stage = std::strcmp(stage, "before") == 0; + const bool after_stage = std::strcmp(stage, "after") == 0; + if (!before_stage && !after_stage) { + return status; + } + const bool capture_blobs = dispatch_trace_blob_capture_enabled(); + if (before_stage && !capture_blobs) { + return status; + } + const std::filesystem::path root(directory_value); + const std::filesystem::path blobs = root / "blobs"; + const std::filesystem::path trace_path = root / "trace.jsonl"; + const std::string kernel_name = dispatch_trace_kernel_name(context, command); + for (size_t binding_index = 0; binding_index < command.kernel.bindings.size(); ++binding_index) { + const PreparedCommandBinding & binding = command.kernel.bindings[binding_index]; + const bool should_capture = before_stage ? resource_access_reads(binding.binding.access) : + resource_access_writes(binding.binding.access); + if (!should_capture) { + continue; + } + std::vector data; + Status download_status = dispatch_trace_download_binding(context, binding, data); + if (!download_status.success()) { + status.append(download_status); + continue; + } + const uint64_t hash = dispatch_trace_hash_bytes(data.data(), data.size()); + std::string blob_relative; + if (capture_blobs) { + std::ostringstream blob_name; + blob_name << stage << "-cmd" << command_index << "-ord" << command.ordinal << "-bind" << binding_index + << "-value" << binding.binding.value.value << "-off" << binding.binding.offset << "-" + << dispatch_trace_hex(hash) << ".bin"; + const std::filesystem::path blob_path = blobs / blob_name.str(); + status.append(dispatch_trace_write_blob(blob_path, data)); + blob_relative = std::filesystem::relative(blob_path, root).string(); + } + try { + std::filesystem::create_directories(root); + std::ofstream out(trace_path, std::ios::app); + if (!out) { + status.log("failed to open dispatch trace %s", trace_path.string().c_str()); + continue; + } + out << "{" + << "\"stage\":\"" << stage << "\"," + << "\"list_kind\":\"" << list_kind << "\"," + << "\"list_size\":" << list_size << ',' + << "\"program_command_count\":" << program_command_count << ',' + << "\"command_index\":" << command_index << ',' + << "\"ordinal\":" << command.ordinal << ',' + << "\"kernel_id\":" << command.kernel.specialization.kernel_id << ',' + << "\"kernel\":\"" << dispatch_trace_json_escape(kernel_name) << "\"," + << "\"binding_index\":" << binding_index << ',' + << "\"binding_name\":\"" << dispatch_trace_json_escape(binding.binding.name) << "\"," + << "\"value\":" << binding.binding.value.value << ',' + << "\"origin\":\"" << command_binding_origin_name(binding.binding.origin) << "\"," + << "\"access\":\"" << resource_access_name(binding.binding.access) << "\"," + << "\"offset\":" << binding.binding.offset << ',' + << "\"length\":" << binding.binding.length << ',' + << "\"ref_offset\":" << binding.ref.offset << ',' + << "\"ref_length\":" << binding.ref.length << ',' + << "\"hash\":\"" << dispatch_trace_hex(hash) << "\""; + if (!blob_relative.empty()) { + out << ",\"blob\":\"" << dispatch_trace_json_escape(blob_relative) << "\""; + } + out << "}\n"; + } catch (const std::exception & error) { + status.log("failed to write dispatch trace %s: %s", trace_path.string().c_str(), error.what()); + } + } + return status; +} + static Status command_program_metadata_context_valid(const CommandProgramExecutionContext & context) { Status status; if (context.target == nullptr) { @@ -272,7 +684,16 @@ static Dispatch build_dispatch(const ResolvedCommand & command) { dispatch.kernel = command.kernel; dispatch.bindings.reserve(command.bindings.size()); for (const ResolvedCommandBinding & binding : command.bindings) { - dispatch.bindings.push_back({ binding.binding.value, binding.binding.offset, binding.binding.length }); + DispatchBinding dispatch_binding; + dispatch_binding.value = binding.binding.value; + dispatch_binding.offset = binding.binding.offset; + dispatch_binding.length = binding.binding.length; + dispatch_binding.layout = binding.binding.layout; + dispatch_binding.source_type = binding.binding.source_type; + dispatch_binding.input_size = binding.binding.input_size; + dispatch_binding.output_size = binding.binding.output_size; + dispatch_binding.source_length = binding.binding.source_length; + dispatch.bindings.push_back(std::move(dispatch_binding)); } return dispatch; } @@ -280,6 +701,15 @@ static Dispatch build_dispatch(const ResolvedCommand & command) { struct GraphValueAccess { bool read = false; bool write = false; + bool native_read = false; + bool transformed_read = false; + bool layout_conflict = false; + std::string layout = kNativeWeightLayout; + ggml_type source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; + size_t source_length = 0; + size_t materialized_length = 0; }; static std::unordered_map collect_graph_value_access(const CommandProgram & commands) { @@ -303,6 +733,38 @@ static std::unordered_map collect_graph_value_access( access.write = true; break; } + if (binding.access == ResourceAccess::Write) { + continue; + } + if (binding.layout == kNativeWeightLayout) { + access.native_read = true; + if (access.transformed_read) { + access.layout_conflict = true; + } + continue; + } + if (binding.offset != 0 || binding.source_length == 0 || binding.length == 0 || + binding.source_type == GGML_TYPE_COUNT || binding.input_size <= 0 || binding.output_size <= 0) { + access.layout_conflict = true; + continue; + } + if (!access.transformed_read) { + access.transformed_read = true; + access.layout = binding.layout; + access.source_type = binding.source_type; + access.input_size = binding.input_size; + access.output_size = binding.output_size; + access.source_length = binding.source_length; + access.materialized_length = binding.length; + } else if (access.layout != binding.layout || access.source_type != binding.source_type || + access.input_size != binding.input_size || access.output_size != binding.output_size || + access.source_length != binding.source_length || + access.materialized_length != binding.length) { + access.layout_conflict = true; + } + if (access.native_read) { + access.layout_conflict = true; + } } } }; @@ -320,21 +782,40 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec materialized.reserve(bindings.bindings().size()); const std::unordered_map access_by_value = collect_graph_value_access(commands); for (const CommandProgramBinding & binding : bindings.bindings()) { - if (!binding.requires_materialization()) { - materialized.push_back(binding); - continue; - } const auto found_access = access_by_value.find(binding.value.value); const GraphValueAccess access = found_access != access_by_value.end() ? found_access->second : GraphValueAccess{}; - if (binding.weight && access.read && !access.write) { + const bool materialize_weight = binding.weight && access.read && !access.write && + (binding.requires_materialization() || access.transformed_read); + if (binding.length == 0 || (!binding.requires_materialization() && !materialize_weight)) { + materialized.push_back(binding); + continue; + } + if (materialize_weight) { + if (access.layout_conflict) { + status.log("weight value %d has conflicting resident layout requests", binding.value.value); + materialized.push_back(binding); + continue; + } HostWeightSource source; source.host_data = binding.host_data; + source.device_buffer = binding.host_data == nullptr ? binding.buffer : nullptr; source.identity = binding.identity; source.generation = binding.generation; source.capacity = binding.capacity; source.offset = binding.offset; - source.length = binding.length; + source.length = access.transformed_read ? access.source_length : binding.length; + source.materialized_length = access.transformed_read ? access.materialized_length : binding.length; + source.layout = access.layout; + source.source_type = access.source_type; + source.input_size = access.input_size; + source.output_size = access.output_size; + if (source.length > binding.length) { + status.log("weight value %d layout source length %zu exceeds runtime length %zu", binding.value.value, + source.length, binding.length); + materialized.push_back(binding); + continue; + } HostWeightAcquireResult resident = context.host_weights->acquire(context.device, context.stream, *context.host_transfers, source); if (!resident.valid()) { @@ -347,12 +828,18 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec device_binding.buffer = resident.lease.buffer(); device_binding.host_data = nullptr; device_binding.offset = 0; - device_binding.capacity = binding.length; + device_binding.length = resident.lease.length(); + device_binding.capacity = resident.lease.length(); materialized.push_back(device_binding); prepared.resident_host_weights.push_back(std::move(resident.lease)); continue; } + if (binding.host_data == nullptr) { + materialized.push_back(binding); + continue; + } + HostStagingBuffer staging; Status allocation_status = allocate_host_staging_buffer(context.device, binding.length, staging); if (!allocation_status.success()) { @@ -365,6 +852,9 @@ static CommandProgramBindings materialize_host_bindings(const CommandProgramExec staging.host_data = static_cast(binding.host_data) + binding.offset; staging.upload = access.read; staging.download = access.write; + if (staging.upload && context.host_buffers != nullptr) { + staging.source_host_buffer = context.host_buffers->find(staging.host_data, staging.length); + } CommandProgramBinding device_binding = binding; device_binding.buffer = staging.buffer; device_binding.host_data = nullptr; @@ -415,8 +905,7 @@ static Status collect_program_constant_images(const CommandProgram & commands, image = &images[found->second]; } - std::copy(initialization.data.begin(), initialization.data.end(), - image->data.begin() + initialization.offset); + std::copy(initialization.data.begin(), initialization.data.end(), image->data.begin() + initialization.offset); } return status; } @@ -438,7 +927,8 @@ static Status validate_program_constant_access(const CommandProgram & } ProgramConstantImage & image = images[found->second]; if (resource_access_writes(binding.access)) { - status.log("constant initialization %s is written by command %u; prepared constant buffer copy " + status.log( + "constant initialization %s is written by command %u; prepared constant buffer copy " "support is required", image.name.c_str(), command.ordinal); continue; @@ -464,8 +954,7 @@ static PreparedProgramConstantBuffer make_program_constant_buffer(ValueId v return result; } -static Status bind_prepared_command_list_program_constants( - const PreparedCommandProgram & prepared, +static Status bind_prepared_command_list_program_constants(const PreparedCommandProgram & prepared, std::vector & prepared_commands) { Status status; for (PreparedCommand & command : prepared_commands) { @@ -526,9 +1015,8 @@ static Status prepare_program_constant_buffers(const CommandProgramExecutionCont continue; } hrx_buffer_t buffer = nullptr; - if (ErrorResult error = take_status( - hrx_allocator_allocate_buffer(hrx_device_allocator(context.device), params, image.data.size(), - &buffer))) { + if (ErrorResult error = take_status(hrx_allocator_allocate_buffer(hrx_device_allocator(context.device), params, + image.data.size(), &buffer))) { status.log("allocate prepared constant %s: %s", image.name.c_str(), error->c_str()); continue; } @@ -557,7 +1045,9 @@ static Status prepare_program_constant_buffers(const CommandProgramExecutionCont return status; } -static Status rebind_prepared_host_staging(const CommandProgramBindings & bindings, PreparedCommandProgram & prepared) { +static Status rebind_prepared_host_staging(const CommandProgramExecutionContext & context, + const CommandProgramBindings & bindings, + PreparedCommandProgram & prepared) { Status status; for (HostStagingBuffer & staging : prepared.host_staging) { const CommandProgramBinding * binding = bindings.find(ValueId(staging.value)); @@ -567,6 +1057,10 @@ static Status rebind_prepared_host_staging(const CommandProgramBindings & bindin continue; } staging.host_data = static_cast(binding->host_data) + binding->offset; + staging.source_host_buffer = HostBufferRef{}; + if (staging.upload && context.host_buffers != nullptr) { + staging.source_host_buffer = context.host_buffers->find(staging.host_data, staging.length); + } } return status; } @@ -581,34 +1075,110 @@ static Status upload_prepared_host_staging(const CommandProgramExecutionContext status.log("missing HRX host transfer manager"); return status; } + const DebugSerialExecutionTrace debug(context); for (const HostStagingBuffer & staging : prepared.host_staging) { if (!staging.upload) { continue; } - Status upload_status = - context.host_transfers->upload_async(context.stream, staging.host_data, staging.buffer, 0, staging.length); + const std::string detail = debug.enabled() ? DebugSerialExecutionTrace::host_staging_detail(staging) : + std::string(); + if (!debug.sync("host-staging-upload-buffer-pre", detail)) { + status.log("debug serial sync before host-staging upload failed for value %d", staging.value); + return status; + } + Status upload_status; + if (staging.source_host_buffer.valid()) { + if (ErrorResult error = take_status(hrx_stream_copy_buffer(context.stream, staging.source_host_buffer.buffer(), + staging.source_host_buffer.offset(), + staging.buffer, 0, staging.length))) { + upload_status.log("HRX host staging buffer copy failed: %s", error->c_str()); + } + } else { + upload_status = + context.host_transfers->upload_async(context.stream, staging.host_data, staging.buffer, 0, staging.length); + } status.append(upload_status); + if (!debug.sync("host-staging-upload-buffer-post", detail)) { + status.log("debug serial sync after host-staging upload failed for value %d", staging.value); + return status; + } + } + return status; +} + +static void mark_graph_replay_synchronized(const CommandProgramExecutionContext & context) { + if (context.graph_replay_state != nullptr) { + context.graph_replay_state->mark_stream_synchronized(); + } +} + +static Status enqueue_stream_execution_barrier(const CommandProgramExecutionContext & context, const char * label) { + Status status; + if (ErrorResult error = take_status(hrx_stream_execution_barrier(context.stream))) { + status.log("%s failed: %s", label, error->c_str()); + } + return status; +} + +static Status flush_stream_commands(const CommandProgramExecutionContext & context, const char * label) { + Status status; + if (ErrorResult error = take_status(hrx_stream_flush(context.stream))) { + status.log("%s failed: %s", label, error->c_str()); + } + return status; +} + +static Status wait_stream_commands(const CommandProgramExecutionContext & context, const char * label) { + Status status; + if (ErrorResult error = take_status(hrx_stream_wait(context.stream))) { + status.log("%s failed: %s", label, error->c_str()); } return status; } static Status download_prepared_host_staging(const CommandProgramExecutionContext & context, - const PreparedCommandProgram & prepared) { + const PreparedCommandProgram & prepared, + bool insert_barrier = true) { Status status; - if (prepared.host_staging.empty()) { + + bool has_download = false; + for (const HostStagingBuffer & staging : prepared.host_staging) { + has_download = has_download || staging.download; + } + if (!has_download) { return status; } - if (context.host_transfers == nullptr) { - status.log("missing HRX host transfer manager"); + if (context.graph_replay_state == nullptr) { + status.log("missing HRX stream completion state"); return status; } + if (insert_barrier) { + Status barrier_status = enqueue_stream_execution_barrier(context, "insert HRX host download barrier"); + if (!barrier_status.success()) { + return barrier_status; + } + } for (const HostStagingBuffer & staging : prepared.host_staging) { if (!staging.download) { continue; } - Status download_status = context.host_transfers->download_synchronous( - context.stream, staging.buffer, 0, staging.host_data, staging.length); - status.append(download_status); + hrx_buffer_t download_buffer = nullptr; + void * download_data = nullptr; + Status allocation_status = + allocate_mapped_host_staging_buffer(context.device, staging.length, download_buffer, download_data); + if (!allocation_status.success()) { + status.log("allocate HRX host download staging buffer for value %d failed", staging.value); + status.append(allocation_status); + return status; + } + if (ErrorResult error = take_status(hrx_stream_copy_buffer(context.stream, staging.buffer, 0, + download_buffer, 0, staging.length))) { + status.log("HRX host download staging copy failed for value %d: %s", staging.value, error->c_str()); + hrx_buffer_release(download_buffer); + return status; + } + context.graph_replay_state->add_host_writeback(staging.host_data, download_data, staging.length, download_buffer); + hrx_buffer_release(download_buffer); } return status; } @@ -649,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; @@ -678,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(), @@ -720,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()); } @@ -759,26 +1330,193 @@ static bool bind_prepared_command_list_transients(const CommandProgram & return true; } -static bool execute_prepared_command_list(const CommandProgramExecutionContext & context, - const std::vector & commands) { - for (const PreparedCommand & command : commands) { +static bool execute_prepared_command_list_serial(const CommandProgramExecutionContext & context, + const std::vector & commands, + const char * list_kind, + size_t program_command_count) { + const DebugSerialExecutionTrace debug(context); + const bool trace_list = debug.enabled() && debug.should_trace_program(program_command_count); + if (trace_list) { + debug.log("begin", "command-list", + DebugSerialExecutionTrace::list_detail(list_kind, commands.size(), program_command_count)); + } + for (size_t i = 0; i < commands.size(); ++i) { + const PreparedCommand & command = commands[i]; + const bool trace_command = debug.enabled() && debug.should_trace_command(program_command_count, i); + const std::string command_detail = trace_command ? + DebugSerialExecutionTrace::command_detail(list_kind, commands.size(), i, + command) : + std::string(); + if (!debug.sync("command-pre-dispatch", command_detail)) { + return false; + } + Status trace_status = + dispatch_trace_command_bindings(context, command, list_kind, commands.size(), i, program_command_count, + "before"); + if (!trace_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(trace_status)); + return false; + } + if (trace_command) { + debug.log("begin", "command-dispatch", command_detail); + } if (!execute_prepared_kernel_command(context, command)) { return false; } + if (trace_command) { + debug.log("end", "command-dispatch", command_detail); + } + if (!debug.sync("command-post-dispatch", command_detail)) { + return false; + } + trace_status = + dispatch_trace_command_bindings(context, command, list_kind, commands.size(), i, program_command_count, + "after"); + if (!trace_status.success()) { + GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(trace_status)); + return false; + } + } + if (trace_list) { + debug.log("end", "command-list", + DebugSerialExecutionTrace::list_detail(list_kind, commands.size(), program_command_count)); } return true; } -struct GraphDependencyChain { - hrx_graph_node_t last = nullptr; +static bool execute_prepared_command_list(const CommandProgramExecutionContext & context, + const std::vector & commands) { + return execute_prepared_command_list_serial(context, commands, "main", commands.size()); +} - const hrx_graph_node_t * deps() const { return last == nullptr ? nullptr : &last; } - size_t dep_count() const { return last == nullptr ? 0 : 1; } - void update(hrx_graph_node_t node) { last = node; } +struct GraphResourceState { + size_t end = 0; + hrx_graph_node_t last_writer = nullptr; + std::vector readers; +}; + +class GraphDependencyPlanner { + public: + std::vector dependencies(const std::vector & bindings) const { + std::vector result; + for (const PreparedCommandBinding & binding : bindings) { + collect_dependencies(binding.ref.buffer, binding.ref.offset, binding.ref.length, + resource_access_writes(binding.binding.access), result); + } + return result; + } + + void record(hrx_graph_node_t node, const std::vector & bindings) { + for (const PreparedCommandBinding & binding : bindings) { + update(node, binding.ref.buffer, binding.ref.offset, binding.ref.length, + resource_access_writes(binding.binding.access)); + } + } + + void record(hrx_graph_node_t node, hrx_buffer_t buffer, size_t offset, size_t length, bool writes) { + update(node, buffer, offset, length, writes); + } + + private: + using ResourceMap = std::map; + + static size_t range_end(size_t offset, size_t length) { + return length > std::numeric_limits::max() - offset ? std::numeric_limits::max() : + offset + length; + } + + static void append_unique(std::vector & nodes, hrx_graph_node_t node) { + if (node != nullptr && std::find(nodes.begin(), nodes.end(), node) == nodes.end()) { + nodes.push_back(node); + } + } + + static void split(ResourceMap & ranges, size_t offset) { + auto upper = ranges.upper_bound(offset); + if (upper == ranges.begin()) { + return; + } + auto current = std::prev(upper); + if (offset <= current->first || offset >= current->second.end) { + return; + } + GraphResourceState right = current->second; + current->second.end = offset; + ranges.emplace(offset, std::move(right)); + } + + void collect_dependencies(hrx_buffer_t buffer, + size_t offset, + size_t length, + bool writes, + std::vector & result) const { + if (buffer == nullptr || length == 0) { + return; + } + const auto resource = resources_.find(buffer); + if (resource == resources_.end()) { + return; + } + const size_t end = range_end(offset, length); + auto range = resource->second.upper_bound(offset); + if (range != resource->second.begin()) { + --range; + if (range->second.end <= offset) { + ++range; + } + } + for (; range != resource->second.end() && range->first < end; ++range) { + append_unique(result, range->second.last_writer); + if (writes) { + for (hrx_graph_node_t reader : range->second.readers) { + append_unique(result, reader); + } + } + } + } + + void update(hrx_graph_node_t node, hrx_buffer_t buffer, size_t offset, size_t length, bool writes) { + if (buffer == nullptr || length == 0) { + return; + } + const size_t end = range_end(offset, length); + ResourceMap & ranges = resources_[buffer]; + split(ranges, offset); + split(ranges, end); + + size_t cursor = offset; + auto range = ranges.lower_bound(offset); + while (cursor < end) { + if (range == ranges.end() || range->first > cursor) { + const size_t gap_end = range == ranges.end() ? end : std::min(end, range->first); + GraphResourceState state; + state.end = gap_end; + if (writes) { + state.last_writer = node; + } else { + state.readers.push_back(node); + } + ranges.emplace(cursor, std::move(state)); + cursor = gap_end; + continue; + } + + if (writes) { + range->second.last_writer = node; + range->second.readers.clear(); + } else { + append_unique(range->second.readers, node); + } + cursor = range->second.end; + ++range; + } + } + + std::unordered_map resources_; }; static Status record_completion_counter_fill(hrx_graph_t graph, - GraphDependencyChain & chain, + GraphDependencyPlanner & dependencies, const CommandProgram & commands, const TransientArenaAllocationRef & allocation) { Status status; @@ -803,17 +1541,17 @@ static Status record_completion_counter_fill(hrx_graph_t sizeof(uint32_t), }; hrx_graph_node_t node = nullptr; - if (ErrorResult error = - take_status(hrx_graph_add_fill_buffer_node(graph, chain.deps(), chain.dep_count(), &attrs, &node))) { + if (ErrorResult error = take_status(hrx_graph_add_fill_buffer_node(graph, nullptr, 0, &attrs, &node))) { status.log("record completion counter fill: %s", error->c_str()); return status; } - chain.update(node); + dependencies.record(node, allocation.buffer, commands.completion_counters.arena_offset, + commands.completion_counters.byte_count, true); return status; } static Status record_prepared_kernel_command(hrx_graph_t graph, - GraphDependencyChain & chain, + GraphDependencyPlanner & dependencies, const PreparedCommand & command) { Status status; const std::string command_context = format_prepared_command_context(command); @@ -845,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(), @@ -857,23 +1595,24 @@ static Status record_prepared_kernel_command(hrx_graph_t graph, refs.size(), 0, }; + const std::vector dependency_nodes = dependencies.dependencies(command.kernel.bindings); hrx_graph_node_t node = nullptr; - if (ErrorResult error = - take_status(hrx_graph_add_kernel_node(graph, chain.deps(), chain.dep_count(), &attrs, &node))) { + if (ErrorResult error = take_status( + hrx_graph_add_kernel_node(graph, dependency_nodes.data(), dependency_nodes.size(), &attrs, &node))) { status.log("record %s: %s", command_context.c_str(), error->c_str()); return status; } - chain.update(node); + dependencies.record(node, command.kernel.bindings); return status; } static Status record_prepared_command_list(hrx_graph_t graph, - GraphDependencyChain & chain, + GraphDependencyPlanner & dependencies, const std::vector & commands, size_t & dispatch_count) { Status status; for (const PreparedCommand & command : commands) { - Status command_status = record_prepared_kernel_command(graph, chain, command); + Status command_status = record_prepared_kernel_command(graph, dependencies, command); if (!command_status.success()) { status.append(command_status); return status; @@ -883,6 +1622,40 @@ static Status record_prepared_command_list(hrx_graph_t gr return status; } +static void retain_prepared_command_list_resources(const std::vector & commands, + std::unordered_map & retained_buffers, + std::unordered_map & retained_executables, + RecordedCommandGraph & recorded) { + for (const PreparedCommand & command : commands) { + if (command.kernel.executable != nullptr && command.kernel.executable->executable != nullptr && + retained_executables.emplace(command.kernel.executable->executable, true).second) { + hrx_executable_retain(command.kernel.executable->executable); + recorded.retained_executables.push_back(command.kernel.executable->executable); + } + for (const PreparedCommandBinding & binding : command.kernel.bindings) { + if (binding.ref.buffer != nullptr && retained_buffers.emplace(binding.ref.buffer, true).second) { + hrx_buffer_retain(binding.ref.buffer); + recorded.retained_buffers.push_back(binding.ref.buffer); + } + } + } +} + +static void retain_prepared_command_graph_resources(const PreparedCommandProgram & prepared, + const TransientArenaAllocationRef & allocation, + RecordedCommandGraph & recorded) { + std::unordered_map retained_buffers; + std::unordered_map retained_executables; + if (allocation.buffer != nullptr) { + hrx_buffer_retain(allocation.buffer); + recorded.retained_buffers.push_back(allocation.buffer); + retained_buffers.emplace(allocation.buffer, true); + } + retain_prepared_command_list_resources(prepared.initialization_commands, retained_buffers, retained_executables, + recorded); + retain_prepared_command_list_resources(prepared.commands, retained_buffers, retained_executables, recorded); +} + static RecordedCommandGraph record_prepared_command_graph(const CommandProgramExecutionContext & context, const CommandProgram & commands, const PreparedCommandProgram & prepared, @@ -900,22 +1673,25 @@ static RecordedCommandGraph record_prepared_command_graph(const CommandProgramEx } recorded.graph = graph; - GraphDependencyChain chain; - recorded.status.append(record_completion_counter_fill(recorded.graph, chain, commands, allocation)); + GraphDependencyPlanner dependencies; + recorded.status.append(record_completion_counter_fill(recorded.graph, dependencies, commands, allocation)); if (!recorded.status.success()) { return recorded; } - recorded.status.append(record_prepared_command_list(recorded.graph, chain, prepared.initialization_commands, + recorded.status.append(record_prepared_command_list(recorded.graph, dependencies, prepared.initialization_commands, recorded.dispatch_count)); if (!recorded.status.success()) { return recorded; } - recorded.status.append(record_prepared_command_list(recorded.graph, chain, prepared.commands, - recorded.dispatch_count)); + recorded.status.append( + record_prepared_command_list(recorded.graph, dependencies, prepared.commands, recorded.dispatch_count)); if (!recorded.status.success()) { return recorded; } + // Unretained HRX command buffers require graph-lifetime resources. + retain_prepared_command_graph_resources(prepared, allocation, recorded); + hrx_graph_exec_t exec = nullptr; if (ErrorResult error = take_status(hrx_graph_instantiate(recorded.graph, 0, &exec))) { recorded.status.log("instantiate HRX graph replay: %s", error->c_str()); @@ -926,8 +1702,83 @@ static RecordedCommandGraph record_prepared_command_graph(const CommandProgramEx return recorded; } +static bool execute_prepared_kernel_command_via_graph(const CommandProgramExecutionContext & context, + const PreparedCommand & command) { + hrx_graph_t graph = nullptr; + if (take_status(hrx_graph_create(context.device, 0, &graph))) { + return false; + } + GraphDependencyPlanner dependencies; + Status status = record_prepared_kernel_command(graph, dependencies, command); + hrx_graph_exec_t exec = nullptr; + if (status.success()) { + if (ErrorResult error = take_status(hrx_graph_instantiate(graph, 0, &exec))) { + status.log("instantiate diagnostic graph: %s", error->c_str()); + } + } + if (status.success()) { + if (ErrorResult error = take_status(hrx_graph_exec_launch(exec, context.stream))) { + status.log("launch diagnostic graph: %s", error->c_str()); + } + } + if (status.success()) { + if (ErrorResult error = take_status(hrx_stream_synchronize(context.stream))) { + status.log("synchronize diagnostic graph: %s", error->c_str()); + } + } + if (exec != nullptr) { + hrx_graph_exec_release(exec); + } + hrx_graph_release(graph); + return status.success(); +} + +static bool execute_prepared_command_prefix_via_graph(const CommandProgramExecutionContext & context, + const std::vector & commands, + size_t prefix_count) { + if (prefix_count == 0 || prefix_count > commands.size()) { + return prefix_count == 0; + } + + hrx_graph_t graph = nullptr; + if (take_status(hrx_graph_create(context.device, 0, &graph))) { + return false; + } + GraphDependencyPlanner dependencies; + Status status; + for (size_t i = 0; i < prefix_count && status.success(); ++i) { + status.append(record_prepared_kernel_command(graph, dependencies, commands[i])); + } + + hrx_graph_exec_t exec = nullptr; + if (status.success()) { + if (ErrorResult error = take_status(hrx_graph_instantiate(graph, 0, &exec))) { + status.log("instantiate diagnostic prefix graph: %s", error->c_str()); + } + } + if (status.success()) { + if (ErrorResult error = take_status(hrx_graph_exec_launch(exec, context.stream))) { + status.log("launch diagnostic prefix graph: %s", error->c_str()); + } + } + if (status.success()) { + if (ErrorResult error = take_status(hrx_stream_synchronize(context.stream))) { + status.log("synchronize diagnostic prefix graph: %s", error->c_str()); + } + } + if (exec != nullptr) { + hrx_graph_exec_release(exec); + } + hrx_graph_release(graph); + return status.success(); +} + } // namespace +bool debug_serial_command_execution_enabled() { + return environment_flag_enabled("GGML_HRX_DEBUG_SERIAL_EXECUTION") || dispatch_trace_enabled(); +} + PreparedCommandProgram prepare_command_program(const CommandProgramExecutionContext & context, const CommandProgram & commands, const CommandProgramBindings & bindings) { @@ -973,7 +1824,6 @@ PreparedCommandProgram prepare_command_program(const CommandProgramExecutionCont prepared.status.append(resolved.status); return prepared; } - std::vector initialization_executable_refs; std::vector command_executable_refs; prepare_command_list(context, resolved.initialization_commands, prepared.initialization_commands, @@ -1022,24 +1872,38 @@ bool bind_and_execute_prepared_command_program(const CommandProgramExecutionCont if (!prepared.valid()) { return execute_prepared_command_program(context, prepared); } - Status rebind_status = rebind_prepared_host_staging(bindings, prepared); + const DebugSerialExecutionTrace debug(context, commands); + debug.log_program_begin(); + HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-rebind-pre"); + Status rebind_status = rebind_prepared_host_staging(context, bindings, prepared); if (!rebind_status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(rebind_status)); return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-rebind-post"); if (commands.transients.arena_size == 0) { + HRX_DEBUG_SERIAL_SYNC(debug, "constant-initialization-pre"); Status status = initialize_command_program_constants(context, commands, {}, prepared); if (!status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "constant-initialization-post"); + HRX_DEBUG_SERIAL_SYNC(debug, "completion-counter-initialization-pre"); status = initialize_command_program_completion_counters(context, commands, {}); if (!status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); return false; } - return bind_prepared_command_program_transients(commands, {}, prepared) && - execute_prepared_command_program(context, prepared); + HRX_DEBUG_SERIAL_SYNC(debug, "completion-counter-initialization-post"); + HRX_DEBUG_SERIAL_SYNC(debug, "transient-bind-pre"); + if (!bind_prepared_command_program_transients(commands, {}, prepared)) { + return false; + } + HRX_DEBUG_SERIAL_SYNC(debug, "transient-bind-post"); + const bool success = execute_prepared_command_program(context, prepared); + debug.log_program_end(success); + return success; } Status status = command_program_transient_context_valid(context, commands); @@ -1048,26 +1912,36 @@ bool bind_and_execute_prepared_command_program(const CommandProgramExecutionCont return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "transient-arena-acquire-pre"); TransientArena::AllocationLease lease = context.transient_arena->acquire_allocation_lease(); status = lease.ensure_capacity(context.device, context.stream, commands.transients.arena_size); if (!status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "transient-arena-acquire-post"); + HRX_DEBUG_SERIAL_SYNC(debug, "transient-bind-pre"); if (!bind_prepared_command_program_transients(commands, lease.current_allocation(), prepared)) { return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "transient-bind-post"); + HRX_DEBUG_SERIAL_SYNC(debug, "constant-initialization-pre"); status = initialize_command_program_constants(context, commands, lease.current_allocation(), prepared); if (!status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "constant-initialization-post"); + HRX_DEBUG_SERIAL_SYNC(debug, "completion-counter-initialization-pre"); status = initialize_command_program_completion_counters(context, commands, lease.current_allocation()); if (!status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(status)); return false; } - return execute_prepared_command_program(context, prepared); + HRX_DEBUG_SERIAL_SYNC(debug, "completion-counter-initialization-post"); + const bool success = execute_prepared_command_program(context, prepared); + debug.log_program_end(success); + return success; } RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( @@ -1092,7 +1966,7 @@ RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( return result; } - Status rebind_status = rebind_prepared_host_staging(bindings, prepared); + Status rebind_status = rebind_prepared_host_staging(context, bindings, prepared); if (!rebind_status.success()) { result.status.append(rebind_status); result.event = HrxGraphReplayEvent::BuildFailed; @@ -1130,11 +2004,79 @@ RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( } const bool had_recorded = recorded.valid(); + + const char * diagnostic_kernel = std::getenv("GGML_HRX_DIAGNOSTIC_GRAPH_KERNEL_ID"); + if (diagnostic_kernel != nullptr && diagnostic_kernel[0] != '\0') { + const uint64_t diagnostic_graph_kernel_id = std::strtoull(diagnostic_kernel, nullptr, 0); + Status diagnostic_upload_status = upload_prepared_host_staging(context, prepared); + if (!diagnostic_upload_status.success()) { + result.status.append(diagnostic_upload_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + auto execute_diagnostic_list = [&](const std::vector & command_list) { + for (const PreparedCommand & command : command_list) { + const bool ok = (diagnostic_graph_kernel_id == 0 || + command.kernel.specialization.kernel_id == diagnostic_graph_kernel_id) ? + execute_prepared_kernel_command_via_graph(context, command) : + execute_prepared_kernel_command(context, command); + if (!ok) { + return false; + } + } + return true; + }; + if (!execute_diagnostic_list(prepared.initialization_commands) || !execute_diagnostic_list(prepared.commands)) { + result.status.log("diagnostic mixed graph execution failed"); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + result.success = true; + return result; + } + + const char * diagnostic_program_size = std::getenv("GGML_HRX_DIAGNOSTIC_GRAPH_PROGRAM_COMMANDS"); + const char * diagnostic_prefix = std::getenv("GGML_HRX_DIAGNOSTIC_GRAPH_PREFIX_COMMANDS"); + if (diagnostic_program_size != nullptr && diagnostic_program_size[0] != '\0' && diagnostic_prefix != nullptr && + diagnostic_prefix[0] != '\0') { + const size_t target_size = std::strtoull(diagnostic_program_size, nullptr, 0); + const size_t prefix_size = std::strtoull(diagnostic_prefix, nullptr, 0); + Status upload_status = upload_prepared_host_staging(context, prepared); + if (!upload_status.success()) { + result.status.append(upload_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + + bool ok = execute_prepared_command_list(context, prepared.initialization_commands); + if (ok && prepared.commands.size() == target_size) { + ok = execute_prepared_command_prefix_via_graph(context, prepared.commands, prefix_size); + for (size_t i = prefix_size; ok && i < prepared.commands.size(); ++i) { + ok = execute_prepared_kernel_command(context, prepared.commands[i]); + } + } else if (ok) { + ok = execute_prepared_command_list(context, prepared.commands); + } + if (!ok) { + result.status.log("diagnostic prefix graph execution failed"); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + Status download_status = download_prepared_host_staging(context, prepared); + if (!download_status.success()) { + result.status.append(download_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + result.success = true; + return result; + } + result.transient_allocation_changed = had_recorded && recorded.bound_transient_arena_allocation_id != prepared.bound_transient_arena_allocation_id; if (!had_recorded || result.transient_allocation_changed) { - result.event = - result.transient_allocation_changed ? HrxGraphReplayEvent::RebuildTransient : HrxGraphReplayEvent::MissBuild; + result.event = result.transient_allocation_changed ? HrxGraphReplayEvent::RebuildTransient : + HrxGraphReplayEvent::MissBuild; const uint64_t build_start_ns = hrx_graph_replay_now_ns(); RecordedCommandGraph rebuilt = record_prepared_command_graph(context, commands, prepared, transient_allocation); result.build_ns = hrx_graph_replay_now_ns() - build_start_ns; @@ -1162,7 +2104,34 @@ RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( result.event = HrxGraphReplayEvent::LaunchFailed; return result; } - Status download_status = download_prepared_host_staging(context, prepared); + Status replay_flush_status = flush_stream_commands(context, "flush HRX graph replay commands"); + if (!replay_flush_status.success()) { + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.status.append(replay_flush_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + Status replay_wait_status = wait_stream_commands(context, "wait for HRX graph replay commands"); + if (!replay_wait_status.success()) { + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.status.append(replay_wait_status); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + const char * diagnostic_sync_size = std::getenv("GGML_HRX_DIAGNOSTIC_GRAPH_SYNC_PROGRAM_COMMANDS"); + const bool diagnostic_sync = std::getenv("GGML_HRX_DIAGNOSTIC_GRAPH_SYNC") != nullptr || + (diagnostic_sync_size != nullptr && diagnostic_sync_size[0] != '\0' && + recorded.dispatch_count == std::strtoull(diagnostic_sync_size, nullptr, 0)); + if (diagnostic_sync) { + if (ErrorResult error = take_status(hrx_stream_synchronize(context.stream))) { + result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; + result.status.log("synchronize diagnostic graph replay: %s", error->c_str()); + result.event = HrxGraphReplayEvent::LaunchFailed; + return result; + } + mark_graph_replay_synchronized(context); + } + Status download_status = download_prepared_host_staging(context, prepared, false); if (!download_status.success()) { result.launch_ns = hrx_graph_replay_now_ns() - launch_start_ns; result.status.append(download_status); @@ -1184,20 +2153,25 @@ bool execute_prepared_command_program(const CommandProgramExecutionContext & con if (!prepared_execution_context_valid(context)) { return false; } + const DebugSerialExecutionTrace debug(context, commands); + HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-upload-pre"); Status upload_status = upload_prepared_host_staging(context, commands); if (!upload_status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(upload_status)); return false; } - if (!execute_prepared_command_list(context, commands.initialization_commands) || - !execute_prepared_command_list(context, commands.commands)) { + HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-upload-post"); + if (!execute_prepared_command_list_serial(context, commands.initialization_commands, "init", commands.commands.size()) || + !execute_prepared_command_list_serial(context, commands.commands, "main", commands.commands.size())) { return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-download-pre"); Status download_status = download_prepared_host_staging(context, commands); if (!download_status.success()) { GGML_LOG_ERROR("%s: %s\n", __func__, status_first_error(download_status)); return false; } + HRX_DEBUG_SERIAL_SYNC(debug, "host-staging-download-post"); return true; } diff --git a/ggml/src/ggml-hrx/runtime/command-program-executor.h b/ggml/src/ggml-hrx/runtime/command-program-executor.h index bfc8689b53b8..bc1debfe03a3 100644 --- a/ggml/src/ggml-hrx/runtime/command-program-executor.h +++ b/ggml/src/ggml-hrx/runtime/command-program-executor.h @@ -4,9 +4,11 @@ #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" +#include #include #include #include @@ -15,6 +17,7 @@ typedef struct hrx_device_s * hrx_device_t; typedef struct hrx_stream_s * hrx_stream_t; typedef struct hrx_buffer_s * hrx_buffer_t; +typedef struct hrx_executable_s * hrx_executable_t; typedef struct hrx_graph_s * hrx_graph_t; typedef struct hrx_graph_exec_s * hrx_graph_exec_t; struct ggml_hrx_loom_jit_amdgpu; @@ -24,6 +27,29 @@ namespace ggml::hrx { class KernelExecutableCache; struct KernelExecutable; class TransientArena; +class HostBufferRegistry; + +struct PendingHostWriteback { + void * host_destination = nullptr; + const void * mapped_source = nullptr; + size_t size = 0; + hrx_buffer_t retained_buffer = nullptr; +}; + +struct GraphReplayStreamState { + GraphReplayStreamState() = default; + ~GraphReplayStreamState(); + + GraphReplayStreamState(const GraphReplayStreamState &) = delete; + GraphReplayStreamState & operator=(const GraphReplayStreamState &) = delete; + + void mark_stream_synchronized(); + void add_host_writeback(void * host_destination, const void * mapped_source, size_t size, hrx_buffer_t buffer); + void clear(); + + private: + std::vector pending_host_writebacks_; +}; struct CommandProgramExecutionContext { hrx_device_t device = nullptr; @@ -34,6 +60,8 @@ struct CommandProgramExecutionContext { TransientArena * transient_arena = nullptr; HostTransferManager * host_transfers = nullptr; HostWeightCache * host_weights = nullptr; + HostBufferRegistry * host_buffers = nullptr; + GraphReplayStreamState * graph_replay_state = nullptr; }; struct PreparedCommandBinding { @@ -44,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; }; @@ -85,6 +114,8 @@ struct PreparedCommandProgram { struct RecordedCommandGraph { hrx_graph_t graph = nullptr; hrx_graph_exec_t exec = nullptr; + std::vector retained_buffers; + std::vector retained_executables; uint64_t bound_transient_arena_allocation_id = kInvalidTransientArenaAllocationId; size_t dispatch_count = 0; Status status; @@ -137,6 +168,8 @@ RecordedCommandGraphExecutionResult bind_and_launch_recorded_command_graph( PreparedCommandProgram & prepared, RecordedCommandGraph & recorded); +bool debug_serial_command_execution_enabled(); + bool execute_command_program(const CommandProgramExecutionContext & context, const CommandProgram & commands, const CommandProgramBindings & bindings); diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index 9ac434ef0a47..fa509c839792 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -62,6 +62,7 @@ CommandProgramBindings GraphExecutor::bind_external_value_buffers(const GraphPro binding.generation = value_binding.generation; binding.capacity = value_binding.capacity; binding.weight = value_binding.weight; + binding.empty_value = ggml_nbytes(external.tensor) == 0; } else { status.log("external value %d is not bound", external.value.value); } @@ -92,8 +93,7 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { return result; } - const bool use_graph_prepared = - !lookup.program->has_prepared_program() || lookup.program->can_use_prepared_fast_path(graph); + const bool use_graph_prepared = lookup.program->can_use_prepared_fast_path(graph); GraphProgramMatch binding_match = std::move(lookup.match); if (use_graph_prepared && lookup.program->has_prepared_program()) { binding_match = lookup.program->match_host_staging_graph(graph); @@ -117,11 +117,13 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { &context_.transient_arena, &context_.host_transfers, &context_.host_weights, + &context_.device->host_buffers, + &context_.graph_replay_state, }; const PreparedCommandProgramCacheExecutionResult execution = use_graph_prepared ? lookup.program->execute_with_result(execution_context, bindings) : context_.prepared_programs.execute_with_result(execution_context, lookup.program->uid(), - lookup.program->command_shape(), + lookup.program->command_shape_hash(), lookup.program->commands(), bindings); if (!execution.success) { result.status.append(execution.status); diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index b22f0818ceb1..98d0cff342b6 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -1,25 +1,86 @@ #include "graph-program-cache.h" +#include "dispatch/command-program-dump.h" #include "dispatch/dispatch-scheduler.h" #include "ggml-impl.h" #include "ggml.h" #include #include +#include #include #include namespace ggml::hrx { namespace { -static bool tensor_metadata_matches(const Value & value, const ggml_tensor * tensor) { - if (tensor == nullptr || value.type != tensor->type || value.element_count != ggml_nelements(tensor) || - value.byte_count != ggml_nbytes(tensor) || value.contiguous != ggml_is_contiguous(tensor)) { +static const ggml_tensor * tensor_storage_root(const ggml_tensor * tensor) { + while (tensor != nullptr && tensor->view_src != nullptr) { + tensor = tensor->view_src; + } + return tensor; +} + +static size_t tensor_storage_offset(const ggml_tensor * tensor) { + return tensor != nullptr && tensor->view_src != nullptr ? tensor->view_offs : 0; +} + +static bool tensor_storage_relative_offset(const ggml_tensor * source, const ggml_tensor * tensor, size_t & offset) { + if (source == nullptr || tensor == nullptr || tensor_storage_root(source) != tensor_storage_root(tensor)) { + return false; + } + const size_t source_offset = tensor_storage_offset(source); + const size_t tensor_offset = tensor_storage_offset(tensor); + if (tensor_offset < source_offset) { return false; } + const size_t relative_offset = tensor_offset - source_offset; + if (relative_offset > ggml_nbytes(source) || ggml_nbytes(tensor) > ggml_nbytes(source) - relative_offset) { + return false; + } + offset = relative_offset; + return true; +} + +static bool tensor_alias_matches(const ValueMap & values, + const Value & value, + const ggml_tensor * tensor, + const std::vector & tensor_by_value) { const bool tensor_alias = tensor->view_src != nullptr; const bool value_alias = value.alias_source.value >= 0; - if (tensor_alias && (!value_alias || value.storage_offset != tensor->view_offs)) { + if (!value_alias) { + return !tensor_alias || value.kind == ValueKind::External; + } + if (!tensor_alias) { + return true; + } + + const Value * source_value = values.find(value.alias_source); + if (source_value == nullptr || value.alias_source.value < 0 || + static_cast(value.alias_source.value) >= tensor_by_value.size()) { + return false; + } + const ggml_tensor * source_tensor = tensor_by_value[static_cast(value.alias_source.value)]; + if (source_tensor == nullptr) { + return tensor_alias && value.storage_offset == tensor->view_offs; + } + size_t relative_offset = 0; + if (!tensor_storage_relative_offset(source_tensor, tensor, relative_offset) || + value.storage_offset < source_value->storage_offset) { + return false; + } + return relative_offset == value.storage_offset - source_value->storage_offset; +} + +static bool tensor_metadata_matches(const ValueMap & values, + const Value & value, + const ggml_tensor * tensor, + const std::vector & tensor_by_value) { + if (tensor == nullptr || value.type != tensor->type || value.element_count != ggml_nelements(tensor) || + value.byte_count != ggml_nbytes(tensor) || value.contiguous != ggml_is_contiguous(tensor)) { + return false; + } + if (!tensor_alias_matches(values, value, tensor, tensor_by_value)) { return false; } for (int i = 0; i < GGML_MAX_DIMS; ++i) { @@ -30,6 +91,33 @@ static bool tensor_metadata_matches(const Value & value, const ggml_tensor * ten return true; } +static std::string format_tensor_metadata(const ggml_tensor * tensor) { + if (tensor == nullptr) { + return "null"; + } + std::ostringstream out; + out << ggml_type_name(tensor->type) << " ne=[" << tensor->ne[0] << ',' << tensor->ne[1] << ',' << tensor->ne[2] + << ',' << tensor->ne[3] << "] nb=[" << tensor->nb[0] << ',' << tensor->nb[1] << ',' << tensor->nb[2] << ',' + << tensor->nb[3] << "] elements=" << ggml_nelements(tensor) << " bytes=" << ggml_nbytes(tensor) + << " contiguous=" << (ggml_is_contiguous(tensor) ? 1 : 0); + if (tensor->view_src != nullptr) { + out << " view_offs=" << tensor->view_offs; + } + return out.str(); +} + +static std::string format_value_metadata(const Value & value) { + std::ostringstream out; + out << (value.kind == ValueKind::External ? "external " : "transient ") << ggml_type_name(value.type) << " ne=[" + << value.ne[0] << ',' << value.ne[1] << ',' << value.ne[2] << ',' << value.ne[3] << "] nb=[" << value.nb[0] + << ',' << value.nb[1] << ',' << value.nb[2] << ',' << value.nb[3] << "] elements=" << value.element_count + << " bytes=" << value.byte_count << " contiguous=" << (value.contiguous ? 1 : 0); + if (value.alias_source.value >= 0) { + out << " storage_offset=" << value.storage_offset; + } + return out.str(); +} + static bool graph_node_params_match(const GraphNode & cached_node, const ggml_tensor * current_node) { if (current_node == nullptr) { return false; @@ -57,11 +145,78 @@ static bool graph_replay_should_fallback(HrxGraphReplayEvent event) { return event == HrxGraphReplayEvent::Ineligible || event == HrxGraphReplayEvent::BuildFailed; } +static void collect_command_graph_values(const std::vector & commands, std::vector & values) { + for (const Command & command : commands) { + for (const CommandBinding & binding : command.bindings) { + if (binding.origin == CommandBindingOrigin::GraphValue && binding.value.value >= 0 && + static_cast(binding.value.value) < values.size()) { + values[static_cast(binding.value.value)] = 1; + } + } + } +} + +static std::vector collect_command_graph_values(const CommandProgram & commands, size_t value_count) { + std::vector values(value_count, 0); + collect_command_graph_values(commands.initialization_commands, values); + collect_command_graph_values(commands.commands, values); + return values; +} + +static bool can_skip_external_binding(const Value & value, + const CommandProgram & commands, + size_t value_count, + std::optional> & command_graph_values) { + if (value.element_count != 0 && value.byte_count != 0) { + return false; + } + if (!command_graph_values.has_value()) { + command_graph_values = collect_command_graph_values(commands, value_count); + } + return (*command_graph_values)[static_cast(value.id.value)] == 0; +} + +class TensorValueIndex { + public: + explicit TensorValueIndex(size_t maximum_size) : maximum_size_(maximum_size) {} + ~TensorValueIndex() { ggml_hash_set_free(&tensors_); } + + TensorValueIndex(const TensorValueIndex &) = delete; + TensorValueIndex & operator=(const TensorValueIndex &) = delete; + + std::optional find(const ggml_tensor * tensor) const { + if (tensors_.size == 0) { + return std::nullopt; + } + const size_t index = ggml_hash_find(&tensors_, tensor); + if (index == GGML_HASHSET_FULL || !ggml_bitset_get(tensors_.used, index)) { + return std::nullopt; + } + return values_[index]; + } + + void emplace(const ggml_tensor * tensor, int32_t value) { + if (tensors_.size == 0) { + GGML_ASSERT(maximum_size_ <= SIZE_MAX / 2); + tensors_ = ggml_hash_set_new(2 * maximum_size_); + values_.resize(tensors_.size); + } + const size_t index = ggml_hash_insert(&tensors_, const_cast(tensor)); + GGML_ASSERT(index != GGML_HASHSET_ALREADY_EXISTS); + values_[index] = value; + } + + private: + size_t maximum_size_; + ggml_hash_set tensors_ = {}; + std::vector values_; +}; + static Status bind_current_value(const ValueMap & values, ValueId expected, const ggml_tensor * tensor, std::vector & tensor_by_value, - std::unordered_map & value_by_tensor, + TensorValueIndex & value_by_tensor, const char * role, size_t node_index) { Status status; @@ -81,21 +236,29 @@ static Status bind_current_value(const ValueMap & return status; } + if (existing_tensor == tensor) { + // Both maps are populated together when a value is first seen. + return status; + } + const auto existing_value = value_by_tensor.find(tensor); - if (existing_value != value_by_tensor.end() && existing_value->second != expected.value) { - status.log("node %zu %s tensor maps to cached values %d and %d", node_index, role, existing_value->second, + if (existing_value.has_value() && *existing_value != expected.value) { + status.log("node %zu %s tensor maps to cached values %d and %d", node_index, role, *existing_value, expected.value); return status; } - if (existing_tensor == nullptr && existing_value == value_by_tensor.end() && - !tensor_metadata_matches(*value, tensor)) { - status.log("node %zu %s value %d metadata does not match current tensor", node_index, role, expected.value); + if (existing_tensor == nullptr && !existing_value.has_value() && + !tensor_metadata_matches(values, *value, tensor, tensor_by_value)) { + status.log("node %zu %s value %d metadata does not match current tensor: cached %s current %s", node_index, + role, expected.value, format_value_metadata(*value).c_str(), format_tensor_metadata(tensor).c_str()); return status; } tensor_by_value[static_cast(expected.value)] = tensor; - value_by_tensor.emplace(tensor, expected.value); + if (!existing_value.has_value()) { + value_by_tensor.emplace(tensor, expected.value); + } return status; } @@ -104,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; } @@ -114,7 +278,9 @@ static std::string command_program_shape_key(const CommandProgram & commands) { out << "|bindings=" << command.bindings.size(); for (const CommandBinding & binding : command.bindings) { out << "|b:" << binding.name << ':' << binding.value.value << ':' << static_cast(binding.origin) << ':' - << binding.offset << ':' << binding.length << ':' << static_cast(binding.access); + << binding.offset << ':' << binding.length << ':' << static_cast(binding.access) << ':' + << binding.layout << ':' << static_cast(binding.source_type) << ':' << binding.input_size << ':' + << binding.output_size << ':' << binding.source_length; } out << "|deps=" << command.dependencies.size(); for (const uint32_t dependency : command.dependencies) { @@ -141,7 +307,8 @@ GraphProgram::GraphProgram(uint64_t uid, target_(std::move(target)), graph_(std::move(graph)), commands_(std::move(commands)), - command_shape_(std::move(command_shape)) {} + command_shape_(std::move(command_shape)), + command_shape_hash_(command_program_shape_hash(command_shape_)) {} const GraphProgramExternalSlot * GraphProgram::find_external_slot(ValueId value) const { const auto found = external_slot_by_value_.find(value.value); @@ -287,7 +454,8 @@ bool GraphProgram::has_prepared_program() const { } bool GraphProgram::can_use_prepared_fast_path(const ggml_cgraph & graph) const { - return graph.nodes == fast_path_nodes_; + // The UID prevents graph-arena address reuse from passing the fast path. + return graph.uid == uid_ && graph.nodes == fast_path_nodes_; } PreparedCommandProgramCacheStats GraphProgram::prepared_stats() const { @@ -322,6 +490,22 @@ PreparedCommandProgramCacheExecutionResult GraphProgram::execute_with_result( ++prepared_stats_.hits; } + if (environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_REPLAY") || debug_serial_command_execution_enabled()) { + result.graph_replay_event = HrxGraphReplayEvent::Disabled; + result.graph_replay_ineligible_reason = debug_serial_command_execution_enabled() ? + "debug_serial_execution" : + "disabled_by_environment"; + result.success = bind_and_execute_prepared_command_program(context, *commands_, bindings, prepared_); + if (!result.success) { + result.status.log("execute cached HRX command program failed"); + } + return result; + } + + if (environment_flag_enabled("GGML_HRX_REBUILD_GRAPH_REPLAY")) { + recorded_ = {}; + } + const RecordedCommandGraphExecutionResult replay = bind_and_launch_recorded_command_graph(context, *commands_, bindings, prepared_, recorded_); apply_graph_replay_result(result, replay); @@ -361,7 +545,7 @@ GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_ const ValueMap & values = graph_->values(); std::vector tensor_by_value(values.size(), nullptr); - std::unordered_map value_by_tensor; + TensorValueIndex value_by_tensor(values.size()); for (size_t node_index = 0; node_index < graph_->nodes().size(); ++node_index) { const GraphNode & cached_node = graph_->nodes()[node_index]; @@ -410,12 +594,23 @@ GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_ } } - for (const ValueId id : values.external_value_ids()) { + std::optional> command_graph_values; + const std::vector external_ids = values.external_value_ids(); + result.external_bindings.reserve(external_ids.size()); + for (const ValueId id : external_ids) { + const Value * value = values.find(id); + if (value == nullptr) { + result.status.log("external value %d is missing from the cached graph", id.value); + return result; + } if (id.value < 0 || static_cast(id.value) >= tensor_by_value.size() || tensor_by_value[static_cast(id.value)] == nullptr) { result.status.log("external value %d is missing from the current graph", id.value); return result; } + if (can_skip_external_binding(*value, *commands_, values.size(), command_graph_values)) { + continue; + } result.external_bindings.push_back({ id, tensor_by_value[static_cast(id.value)] }); } return result; @@ -478,6 +673,7 @@ GraphProgramLookup GraphProgramCache::build_from_imported(const ggml_cgraph & g GraphProgram * cached_program = program.get(); { std::lock_guard lock(mutex_); + validated_matches_.clear(); programs_[graph.uid] = std::move(program); cached_program = programs_[graph.uid].get(); last_program_ = cached_program; @@ -495,6 +691,7 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, if (graph.uid != 0) { const bool disable_fast_path = environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_UID_FAST_PATH"); const bool validate_fast_path = environment_flag_enabled("GGML_HRX_VALIDATE_GRAPH_UID_CACHE"); + const bool allow_structural_reuse = std::getenv("GGML_HRX_DUMP_COMMAND_PROGRAM_DIR") == nullptr; GraphProgram * cached_program = nullptr; { std::lock_guard lock(mutex_); @@ -507,6 +704,13 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, cached_program = found->second.get(); last_program_ = cached_program; } + if (cached_program == nullptr && !disable_fast_path && allow_structural_reuse) { + const auto matched = validated_matches_.find(graph.uid); + if (matched != validated_matches_.end() && matched->second.nodes == graph.nodes && + matched->second.program->target() == target) { + cached_program = matched->second.program; + } + } } } if (cached_program != nullptr) { @@ -535,6 +739,34 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, result.status.append(match.status); return result; } + std::lock_guard lock(mutex_); + validated_matches_.erase(graph.uid); + } + + // Physical bindings have a separate prepared-program cache key. + if (allow_structural_reuse) { + std::lock_guard lock(mutex_); + for (const auto & entry : programs_) { + GraphProgram * candidate = entry.second.get(); + if (candidate == cached_program || candidate->target() != target) { + continue; + } + GraphProgramMatch match = candidate->match_current_graph(graph); + if (!match.valid()) { + continue; + } + last_program_ = candidate; + if (!disable_fast_path) { + if (validated_matches_.size() >= 128) { + validated_matches_.clear(); + } + validated_matches_[graph.uid] = { candidate, graph.nodes }; + } + ++stats_.hits; + result.program = candidate; + result.match = std::move(match); + return result; + } } } @@ -559,6 +791,7 @@ GraphProgramCacheStats GraphProgramCache::stats() const { void GraphProgramCache::clear() { std::lock_guard lock(mutex_); + validated_matches_.clear(); programs_.clear(); last_program_ = nullptr; } @@ -581,6 +814,10 @@ std::unique_ptr GraphProgramCache::build_program_from_imported(uin } std::string command_shape = command_program_shape_key(commands); + Status dump_status = dump_command_program_kernels_if_requested(commands, corpus, target, command_shape); + for (const std::string & error : dump_status.errors()) { + GGML_LOG_WARN("ggml_hrx: %s\n", error.c_str()); + } return std::make_unique(uid, target, std::make_unique(std::move(imported_graph)), std::make_unique(std::move(commands)), std::move(command_shape)); diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index 0d2a804c295c..4cdb3774370a 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -56,6 +56,8 @@ class GraphProgram { const std::string & command_shape() const { return command_shape_; } + uint64_t command_shape_hash() const { return command_shape_hash_; } + const Graph & graph() const { return *graph_; } Graph & graph() { return *graph_; } @@ -89,6 +91,7 @@ class GraphProgram { std::unique_ptr graph_; std::unique_ptr commands_; std::string command_shape_; + uint64_t command_shape_hash_ = 0; std::vector external_slots_; std::unordered_map external_slot_by_value_; @@ -138,6 +141,11 @@ class GraphProgramCache { void clear(); private: + struct ValidatedGraphMatch { + GraphProgram * program = nullptr; + const ggml_tensor * const * nodes = nullptr; + }; + GraphProgramLookup build_from_imported(const ggml_cgraph & graph, Graph && imported_graph, const KernelCorpus & corpus, @@ -151,6 +159,7 @@ class GraphProgramCache { mutable std::mutex mutex_; std::unordered_map> programs_; + std::unordered_map validated_matches_; GraphProgram * last_program_ = nullptr; GraphProgramCacheStats stats_; }; diff --git a/ggml/src/ggml-hrx/runtime/host-memory.cpp b/ggml/src/ggml-hrx/runtime/host-memory.cpp index e74ae16c98d3..c1741fd802ec 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.cpp +++ b/ggml/src/ggml-hrx/runtime/host-memory.cpp @@ -1,15 +1,1206 @@ #include "host-memory.h" +#include "ggml-quants.h" #include "hrx-interop-utils.h" #include "hrx_runtime.h" +#include +#include +#include +#include +#include +#include +#include #include +#include #include +#include namespace ggml::hrx { namespace { static constexpr size_t kMaxInlineUploadBytes = 63 * 1024; +static constexpr size_t kLargeHostUploadBytes = 1024 * 1024; + +static bool mapped_download_staging_enabled() { + static const bool enabled = [] { + const char * mode = std::getenv("GGML_HRX_WEIGHT_DOWNLOAD_STAGING"); + return mode == nullptr || mode[0] == '\0' || std::strcmp(mode, "mapped") == 0; + }(); + return enabled; +} + +struct SymmetricI4Block { + std::array scales = {}; + std::array, 8> payloads = {}; +}; + +static bool checked_multiply(size_t lhs, size_t rhs, size_t & result) { + if (rhs != 0 && lhs > std::numeric_limits::max() / rhs) { + return false; + } + result = lhs * rhs; + return true; +} + +template static bool run_worker_threads(size_t thread_count, Work && work) { + std::atomic worker_failed = false; + std::vector workers; + auto join_workers = [&]() { + for (std::thread & worker : workers) { + if (worker.joinable()) { + worker.join(); + } + } + }; + + try { + workers.reserve(thread_count); + for (size_t thread = 0; thread < thread_count; ++thread) { + workers.emplace_back([&, thread]() { + try { + work(thread); + } catch (...) { + worker_failed.store(true, std::memory_order_relaxed); + } + }); + } + } catch (...) { + join_workers(); + return false; + } + join_workers(); + return !worker_failed.load(std::memory_order_relaxed); +} + +static void dequantize_block(ggml_type type, const uint8_t * source, float * destination) { + switch (type) { + case GGML_TYPE_Q4_K: + dequantize_row_q4_K(reinterpret_cast(source), destination, QK_K); + break; + case GGML_TYPE_Q5_K: + dequantize_row_q5_K(reinterpret_cast(source), destination, QK_K); + break; + case GGML_TYPE_IQ4_XS: + dequantize_row_iq4_xs(reinterpret_cast(source), destination, QK_K); + break; + case GGML_TYPE_Q6_K: + dequantize_row_q6_K(reinterpret_cast(source), destination, QK_K); + break; + default: + break; + } +} + +static bool quantize_symmetric_value(float value, float scale, int quant_min, int quant_max, int & quantized) { + if (!std::isfinite(value) || !std::isfinite(scale) || scale <= 0.0f) { + return false; + } + float quotient = value / scale; + if (std::isnan(quotient)) { + return false; + } + quotient = std::clamp(quotient, static_cast(quant_min), static_cast(quant_max)); + quantized = static_cast(std::nearbyint(quotient)); + return true; +} + +static bool fit_symmetric_scale(const float * values, + size_t value_count, + int quant_min, + int quant_max, + float & encoded_scale) { + float positive_max = 0.0f; + float negative_max = 0.0f; + bool nonzero = false; + for (size_t element = 0; element < value_count; ++element) { + if (!std::isfinite(values[element])) { + return false; + } + nonzero = nonzero || values[element] != 0.0f; + positive_max = std::max(positive_max, values[element]); + negative_max = std::max(negative_max, -values[element]); + } + if (!nonzero) { + encoded_scale = 0.0f; + return true; + } + + float scale = std::max(positive_max / quant_max, negative_max / -quant_min); + if (!std::isfinite(scale) || scale <= 0.0f) { + return false; + } + for (int iteration = 0; iteration < 2; ++iteration) { + double numerator = 0.0; + double denominator = 0.0; + for (size_t element = 0; element < value_count; ++element) { + int quantized = 0; + if (!quantize_symmetric_value(values[element], scale, quant_min, quant_max, quantized)) { + return false; + } + numerator += static_cast(values[element]) * quantized; + denominator += static_cast(quantized) * quantized; + } + if (denominator > 0.0) { + scale = static_cast(numerator / denominator); + if (!std::isfinite(scale) || scale <= 0.0f) { + return false; + } + } + } + + encoded_scale = ggml_fp16_to_fp32(ggml_fp32_to_fp16(scale)); + return std::isfinite(encoded_scale) && encoded_scale > 0.0f; +} + +static bool fit_shared_symmetric_scale(const std::array, 4> & values, + size_t row_count, + size_t group, + int quant_min, + int quant_max, + bool multistart, + float & encoded_scale) { + float positive_max = 0.0f; + float negative_max = 0.0f; + bool nonzero = false; + for (size_t row = 0; row < row_count; ++row) { + const float * group_values = values[row].data() + group * 32; + for (size_t element = 0; element < 32; ++element) { + if (!std::isfinite(group_values[element])) { + return false; + } + nonzero = nonzero || group_values[element] != 0.0f; + positive_max = std::max(positive_max, group_values[element]); + negative_max = std::max(negative_max, -group_values[element]); + } + } + if (!nonzero) { + encoded_scale = 0.0f; + return true; + } + + const float initial_scale = std::max(positive_max / quant_max, negative_max / -quant_min); + if (!std::isfinite(initial_scale) || initial_scale <= 0.0f) { + return false; + } + + auto refit = [&](float & scale) { + const int iteration_count = multistart ? 4 : 2; + for (int iteration = 0; iteration < iteration_count; ++iteration) { + double numerator = 0.0; + double denominator = 0.0; + for (size_t row = 0; row < row_count; ++row) { + const float * group_values = values[row].data() + group * 32; + for (size_t element = 0; element < 32; ++element) { + int quantized = 0; + if (!quantize_symmetric_value(group_values[element], scale, quant_min, quant_max, quantized)) { + return false; + } + numerator += static_cast(group_values[element]) * quantized; + denominator += static_cast(quantized) * quantized; + } + } + if (denominator > 0.0) { + scale = static_cast(numerator / denominator); + if (!std::isfinite(scale) || scale <= 0.0f) { + return false; + } + } + } + return true; + }; + + if (!multistart) { + float scale = initial_scale; + if (!refit(scale)) { + return false; + } + encoded_scale = ggml_fp16_to_fp32(ggml_fp32_to_fp16(scale)); + return std::isfinite(encoded_scale) && encoded_scale > 0.0f; + } + + constexpr std::array kInitialScaleRatios = { 0.50f, 0.60f, 0.70f, 0.80f, 0.90f, 1.00f, 1.10f }; + double best_error = std::numeric_limits::infinity(); + float best_scale = 0.0f; + for (float ratio : kInitialScaleRatios) { + float scale = initial_scale * ratio; + if (!refit(scale)) { + return false; + } + scale = ggml_fp16_to_fp32(ggml_fp32_to_fp16(scale)); + double error = 0.0; + for (size_t row = 0; row < row_count; ++row) { + const float * group_values = values[row].data() + group * 32; + for (size_t element = 0; element < 32; ++element) { + int quantized = 0; + if (!quantize_symmetric_value(group_values[element], scale, quant_min, quant_max, quantized)) { + return false; + } + const double delta = + static_cast(group_values[element]) - static_cast(quantized) * scale; + error += delta * delta; + } + } + if (error < best_error) { + best_error = error; + best_scale = scale; + } + } + + encoded_scale = best_scale; + return std::isfinite(encoded_scale) && encoded_scale > 0.0f; +} + +static bool quantize_symmetric_i4_k64(const float * values, SymmetricI4Block & output) { + for (size_t pair = 0; pair < output.scales.size() / 2; ++pair) { + const float * pair_values = values + pair * 64; + float scale = 0.0f; + if (!fit_symmetric_scale(pair_values, 64, -8, 7, scale)) { + return false; + } + if (scale == 0.0f) { + continue; + } + + const ggml_fp16_t encoded_scale = ggml_fp32_to_fp16(scale); + output.scales[pair * 2] = encoded_scale; + output.scales[pair * 2 + 1] = encoded_scale; + for (size_t half = 0; half < 2; ++half) { + const size_t logical_group = pair * 2 + half; + const float * group_values = pair_values + half * 32; + for (size_t element_pair = 0; element_pair < 16; ++element_pair) { + int low = 0; + int high = 0; + if (!quantize_symmetric_value(group_values[element_pair * 2], scale, -8, 7, low) || + !quantize_symmetric_value(group_values[element_pair * 2 + 1], scale, -8, 7, high)) { + return false; + } + output.payloads[logical_group][element_pair] = + static_cast((low & 0x0F) | ((high & 0x0F) << 4)); + } + } + } + return true; +} + +static bool quantize_symmetric_i4_k32(const float * values, SymmetricI4Block & output) { + for (size_t group = 0; group < output.scales.size(); ++group) { + const float * group_values = values + group * 32; + float scale = 0.0f; + if (!fit_symmetric_scale(group_values, 32, -8, 7, scale)) { + return false; + } + if (scale == 0.0f) { + continue; + } + output.scales[group] = ggml_fp32_to_fp16(scale); + for (size_t element_pair = 0; element_pair < 16; ++element_pair) { + int low = 0; + int high = 0; + if (!quantize_symmetric_value(group_values[element_pair * 2], scale, -8, 7, low) || + !quantize_symmetric_value(group_values[element_pair * 2 + 1], scale, -8, 7, high)) { + return false; + } + output.payloads[group][element_pair] = static_cast((low & 0x0F) | ((high & 0x0F) << 4)); + } + } + return true; +} + +using SymmetricI4Quantize = bool (*)(const float *, SymmetricI4Block &); + +static Status materialize_symmetric_i4_row64(const HostWeightSource & source, + SymmetricI4Quantize quantize, + std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q4_K && source.source_type != GGML_TYPE_Q5_K && + source.source_type != GGML_TYPE_Q6_K && source.source_type != GGML_TYPE_IQ4_XS) { + status.log("layout %s does not support GGML type %d", source.layout.c_str(), + static_cast(source.source_type)); + return status; + } + if (source.input_size <= 0 || source.input_size % QK_K != 0 || source.output_size <= 0 || + source.output_size % 64 != 0) { + status.log("layout %s requires K divisible by %d and rows divisible by 64, got K=%lld rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.input_size), + static_cast(source.output_size)); + return status; + } + + const size_t row_bytes = ggml_row_size(source.source_type, source.input_size); + size_t expected_source_bytes = 0; + if (!checked_multiply(row_bytes, static_cast(source.output_size), expected_source_bytes) || + source.length != expected_source_bytes) { + status.log("layout %s source length %zu does not match expected %zu", source.layout.c_str(), source.length, + expected_source_bytes); + return status; + } + + constexpr size_t materialized_row_bytes = 144; + constexpr size_t row_group = 64; + const size_t logical_row_count = static_cast(source.output_size); + const size_t padded_row_count = (logical_row_count + row_group - 1) / row_group * row_group; + const size_t block_count = static_cast(source.input_size / QK_K); + size_t expected_output_bytes = 0; + if (!checked_multiply(padded_row_count, block_count, expected_output_bytes) || + !checked_multiply(expected_output_bytes, materialized_row_bytes, expected_output_bytes) || + source.materialized_length != expected_output_bytes) { + status.log("layout %s materialized length %zu does not match expected %zu", source.layout.c_str(), + source.materialized_length, expected_output_bytes); + return status; + } + + output.assign(expected_output_bytes, 0); + const auto * source_bytes = static_cast(source.host_data) + source.offset; + constexpr size_t field_count = 9; + constexpr size_t field_bytes = 16; + constexpr size_t block_out_bytes = field_count * field_bytes; + const size_t source_block_bytes = ggml_type_size(source.source_type); + const size_t thread_count = std::min(logical_row_count, std::max(1u, std::thread::hardware_concurrency())); + std::atomic conversion_failed = false; + + auto convert_rows = [&](size_t row_begin, size_t row_end) { + std::array decoded = {}; + for (size_t row = row_begin; row < row_end; ++row) { + if (conversion_failed.load(std::memory_order_relaxed)) { + return; + } + const size_t group = row / row_group; + const size_t lane = row % row_group; + const uint8_t * source_row = source_bytes + row * row_bytes; + for (size_t block = 0; block < block_count; ++block) { + dequantize_block(source.source_type, source_row + block * source_block_bytes, decoded.data()); + SymmetricI4Block converted; + if (!quantize(decoded.data(), converted)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + const size_t group_base = (group * block_count + block) * row_group * block_out_bytes; + uint8_t * header = output.data() + group_base + lane * field_bytes; + for (size_t group = 0; group < converted.scales.size(); ++group) { + std::memcpy(header + group * sizeof(ggml_fp16_t), &converted.scales[group], sizeof(ggml_fp16_t)); + } + for (size_t logical_group = 0; logical_group < converted.payloads.size(); ++logical_group) { + uint8_t * payload = + output.data() + group_base + (1 + logical_group) * row_group * field_bytes + lane * field_bytes; + std::memcpy(payload, converted.payloads[logical_group].data(), field_bytes); + } + } + } + }; + + if (!run_worker_threads(thread_count, [&](size_t thread) { + const size_t row_begin = logical_row_count * thread / thread_count; + const size_t row_end = logical_row_count * (thread + 1) / thread_count; + convert_rows(row_begin, row_end); + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + return status; + } + if (conversion_failed.load(std::memory_order_relaxed)) { + output.clear(); + status.log("layout %s cannot represent non-finite or out-of-range symmetric weights", source.layout.c_str()); + } + return status; +} + +static Status materialize_symmetric_i4_k32_row64(const HostWeightSource & source, std::vector & output) { + return materialize_symmetric_i4_row64(source, quantize_symmetric_i4_k32, output); +} + +static Status materialize_symmetric_i4_k64_row64(const HostWeightSource & source, std::vector & output) { + return materialize_symmetric_i4_row64(source, quantize_symmetric_i4_k64, output); +} + +static Status materialize_symmetric_k32_eightgroups_shared4(const HostWeightSource & source, + int quant_bits, + std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q4_K && source.source_type != GGML_TYPE_Q5_K && + source.source_type != GGML_TYPE_Q6_K && source.source_type != GGML_TYPE_IQ4_XS) { + status.log("layout %s does not support GGML type %d", source.layout.c_str(), + static_cast(source.source_type)); + return status; + } + if (source.input_size <= 0 || source.input_size % QK_K != 0 || source.output_size <= 0) { + status.log("layout %s requires K divisible by %d and positive rows, got K=%lld rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.input_size), + static_cast(source.output_size)); + return status; + } + + constexpr size_t shared_rows = 4; + const size_t logical_rows = static_cast(source.output_size); + if (logical_rows > std::numeric_limits::max() - 255) { + status.log("layout %s row count overflows row-group sizing", source.layout.c_str()); + return status; + } + const size_t materialized_row_bytes = 4 + 8 * 32 * static_cast(quant_bits) / 8; + const size_t unpadded_bytes = + logical_rows * static_cast(source.input_size / QK_K) * materialized_row_bytes; + const size_t dynamic_row_group = ((logical_rows + 255) / 256) * 32; + const size_t row_group = + quant_bits == 4 && unpadded_bytes < size_t{ 16 } * 1024 * 1024 ? 32 : + quant_bits == 4 && unpadded_bytes <= size_t{ 32 } * 1024 * 1024 && source.output_size > source.input_size ? 96 : + dynamic_row_group; + if (logical_rows > std::numeric_limits::max() - (row_group - 1)) { + status.log("layout %s row count overflows row-group padding", source.layout.c_str()); + return status; + } + const size_t physical_rows = (logical_rows + row_group - 1) / row_group * row_group; + + const size_t row_bytes = ggml_row_size(source.source_type, source.input_size); + size_t expected_source_bytes = 0; + if (!checked_multiply(row_bytes, logical_rows, expected_source_bytes) || source.length != expected_source_bytes) { + status.log("layout %s source length %zu does not match expected %zu", source.layout.c_str(), source.length, + expected_source_bytes); + return status; + } + + if (quant_bits != 2 && quant_bits != 4) { + status.log("layout %s requires 2-bit or 4-bit symmetric materialization", source.layout.c_str()); + return status; + } + + const int quant_min = -(1 << (quant_bits - 1)); + const int quant_max = (1 << (quant_bits - 1)) - 1; + const size_t field_bytes = 32 * static_cast(quant_bits) / 8; + const size_t encoded_row_bytes = 4 + 8 * field_bytes; + const size_t block_count = static_cast(source.input_size / QK_K); + size_t expected_output_bytes = 0; + if (!checked_multiply(physical_rows, block_count, expected_output_bytes) || + !checked_multiply(expected_output_bytes, encoded_row_bytes, expected_output_bytes) || + source.materialized_length != expected_output_bytes) { + status.log("layout %s materialized length %zu does not match expected %zu", source.layout.c_str(), + source.materialized_length, expected_output_bytes); + return status; + } + + output.assign(expected_output_bytes, uint8_t{ 0 }); + const auto * source_bytes = static_cast(source.host_data) + source.offset; + const size_t scale_plane_bytes = row_group / shared_rows * 16; + const size_t payload_plane_bytes = row_group * field_bytes; + const size_t payload_block_bytes = 8 * payload_plane_bytes; + const size_t block_out_bytes = scale_plane_bytes + payload_block_bytes; + const size_t row_group_bytes = block_count * block_out_bytes; + const size_t source_block_bytes = ggml_type_size(source.source_type); + const size_t logical_cohorts = (logical_rows + shared_rows - 1) / shared_rows; + const size_t thread_count = std::min(logical_cohorts, std::max(1u, std::thread::hardware_concurrency())); + std::atomic conversion_failed = false; + + auto convert_cohorts = [&](size_t cohort_begin, size_t cohort_end) { + std::array, shared_rows> decoded = {}; + for (size_t cohort = cohort_begin; cohort < cohort_end; ++cohort) { + if (conversion_failed.load(std::memory_order_relaxed)) { + return; + } + const size_t first_row = cohort * shared_rows; + const size_t row_count = std::min(shared_rows, logical_rows - first_row); + const size_t row_group_id = first_row / row_group; + const size_t cohort_lane = first_row % row_group / shared_rows; + for (size_t block = 0; block < block_count; ++block) { + for (size_t row = 0; row < row_count; ++row) { + const uint8_t * source_row = source_bytes + (first_row + row) * row_bytes; + dequantize_block(source.source_type, source_row + block * source_block_bytes, decoded[row].data()); + } + + const size_t group_base = row_group_id * row_group_bytes; + const size_t payload_block_base = group_base + block * payload_block_bytes; + const size_t scale_region = group_base + block_count * payload_block_bytes; + uint8_t * header = output.data() + scale_region + block * scale_plane_bytes + cohort_lane * 16; + std::array scales = {}; + for (size_t group = 0; group < scales.size(); ++group) { + float scale = 0.0f; + const bool multistart = source.layout == kSymmetricI4K32EightGroupsShared4MultistartLayout; + if (!fit_shared_symmetric_scale(decoded, row_count, group, quant_min, quant_max, multistart, + scale)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + const ggml_fp16_t encoded_scale = ggml_fp32_to_fp16(scale); + std::memcpy(header + group * sizeof(encoded_scale), &encoded_scale, sizeof(encoded_scale)); + scales[group] = scale; + } + + for (size_t row = 0; row < row_count; ++row) { + const size_t row_lane = (first_row + row) % row_group; + for (size_t group = 0; group < scales.size(); ++group) { + uint8_t * payload = + output.data() + payload_block_base + group * payload_plane_bytes + row_lane * field_bytes; + const float * group_values = decoded[row].data() + group * 32; + const float scale = scales[group]; + if (scale == 0.0f) { + continue; + } + if (quant_bits == 4) { + for (size_t element_pair = 0; element_pair < 16; ++element_pair) { + int low = 0; + int high = 0; + if (!quantize_symmetric_value(group_values[element_pair * 2], scale, quant_min, + quant_max, low) || + !quantize_symmetric_value(group_values[element_pair * 2 + 1], scale, quant_min, + quant_max, high)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + payload[element_pair] = static_cast((low & 0x0F) | ((high & 0x0F) << 4)); + } + } else { + for (size_t element_quartet = 0; element_quartet < 8; ++element_quartet) { + uint8_t packed = 0; + for (size_t element = 0; element < 4; ++element) { + const size_t index = element_quartet * 4 + element; + int quantized = 0; + if (!quantize_symmetric_value(group_values[index], scale, quant_min, quant_max, + quantized)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + packed |= static_cast((quantized & 0x03) << (element * 2)); + } + payload[element_quartet] = packed; + } + } + } + } + } + } + }; + + if (!run_worker_threads(thread_count, [&](size_t thread) { + const size_t cohort_begin = logical_cohorts * thread / thread_count; + const size_t cohort_end = logical_cohorts * (thread + 1) / thread_count; + convert_cohorts(cohort_begin, cohort_end); + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + return status; + } + if (conversion_failed.load(std::memory_order_relaxed)) { + output.clear(); + status.log("layout %s cannot represent non-finite or out-of-range symmetric weights", source.layout.c_str()); + } + return status; +} + +static Status materialize_symmetric_i4_k32_eightgroups_shared4(const HostWeightSource & source, + std::vector & output) { + return materialize_symmetric_k32_eightgroups_shared4(source, 4, output); +} + +static Status materialize_symmetric_i5_k32(const HostWeightSource & source, std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q5_K) { + status.log("layout %s does not support GGML type %d", source.layout.c_str(), + static_cast(source.source_type)); + return status; + } + if (source.input_size <= 0 || source.input_size % QK_K != 0 || source.output_size <= 0 || + source.output_size % 64 != 0) { + status.log("layout %s requires K divisible by %d and rows divisible by 64, got K=%lld rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.input_size), + static_cast(source.output_size)); + return status; + } + + const size_t row_bytes = ggml_row_size(GGML_TYPE_Q5_K, source.input_size); + const size_t row_count = static_cast(source.output_size); + const size_t block_count = static_cast(source.input_size / QK_K); + size_t expected_bytes = 0; + if (!checked_multiply(row_bytes, row_count, expected_bytes) || source.length != expected_bytes || + source.materialized_length != expected_bytes) { + status.log("layout %s has inconsistent source/materialized lengths for K=%lld rows=%lld", source.layout.c_str(), + static_cast(source.input_size), static_cast(source.output_size)); + return status; + } + + output.assign(expected_bytes, uint8_t{ 0 }); + const auto * source_bytes = static_cast(source.host_data) + source.offset; + const size_t thread_count = + std::min(row_count, static_cast(std::max(1u, std::thread::hardware_concurrency()))); + std::atomic conversion_failed = false; + + auto convert_rows = [&](size_t row_begin, size_t row_end) { + std::array values = {}; + std::array, 8> codes = {}; + for (size_t row = row_begin; row < row_end; ++row) { + const uint8_t * source_row = source_bytes + row * row_bytes; + uint8_t * output_row = output.data() + row * row_bytes; + for (size_t block = 0; block < block_count; ++block) { + dequantize_row_q5_K(reinterpret_cast(source_row) + block, values.data(), QK_K); + uint8_t * record = output_row + block * sizeof(block_q5_K); + for (size_t group = 0; group < 8; ++group) { + const float * group_values = values.data() + group * 32; + float positive_max = 0.0f; + float negative_max = 0.0f; + for (size_t element = 0; element < 32; ++element) { + const float value = group_values[element]; + if (!std::isfinite(value)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + positive_max = std::max(positive_max, value); + negative_max = std::max(negative_max, -value); + } + float scale = std::max(positive_max / 15.0f, negative_max / 16.0f); + if (scale == 0.0f) { + scale = 1.0f; + } + for (int iteration = 0; iteration < 2; ++iteration) { + double numerator = 0.0; + double denominator = 0.0; + for (size_t element = 0; element < 32; ++element) { + int quantized = 0; + if (!quantize_symmetric_value(group_values[element], scale, -16, 15, quantized)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + numerator += static_cast(group_values[element]) * quantized; + denominator += static_cast(quantized) * quantized; + } + if (denominator > 0.0) { + scale = static_cast(numerator / denominator); + } + if (!std::isfinite(scale) || scale <= 0.0f) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + } + + const ggml_fp16_t encoded_scale = ggml_fp32_to_fp16(scale); + std::memcpy(record + group * sizeof(encoded_scale), &encoded_scale, sizeof(encoded_scale)); + scale = ggml_fp16_to_fp32(encoded_scale); + for (size_t element = 0; element < 32; ++element) { + int quantized = 0; + if (!quantize_symmetric_value(group_values[element], scale, -16, 15, quantized)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + codes[group][element] = static_cast(quantized) & uint8_t{ 31 }; + } + } + + uint8_t * high = record + 16; + uint8_t * packed = record + 48; + for (size_t element = 0; element < 32; ++element) { + uint8_t high_byte = 0; + for (size_t group = 0; group < 8; ++group) { + high_byte |= static_cast(((codes[group][element] >> 4) & 1) << group); + } + high[element] = high_byte; + } + for (size_t pair = 0; pair < 4; ++pair) { + const size_t group0 = pair * 2; + const size_t group1 = group0 + 1; + for (size_t element = 0; element < 32; ++element) { + packed[pair * 32 + element] = + static_cast((codes[group0][element] & 15) | ((codes[group1][element] & 15) << 4)); + } + } + } + } + }; + + if (!run_worker_threads(thread_count, [&](size_t thread) { + convert_rows(row_count * thread / thread_count, row_count * (thread + 1) / thread_count); + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + return status; + } + if (conversion_failed.load(std::memory_order_relaxed)) { + output.clear(); + status.log("layout %s cannot represent non-finite symmetric weights", source.layout.c_str()); + } + return status; +} + +static Status materialize_symmetric_i8_k256_row64(const HostWeightSource & source, std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q5_K) { + status.log("layout %s does not support GGML type %d", source.layout.c_str(), + static_cast(source.source_type)); + return status; + } + if (source.input_size <= 0 || source.input_size % QK_K != 0 || source.output_size <= 0 || + source.output_size % 64 != 0) { + status.log("layout %s requires K divisible by %d and rows divisible by 64, got K=%lld rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.input_size), + static_cast(source.output_size)); + return status; + } + + const size_t row_bytes = ggml_row_size(GGML_TYPE_Q5_K, source.input_size); + const size_t row_count = static_cast(source.output_size); + const size_t block_count = static_cast(source.input_size / QK_K); + constexpr size_t row_group = 64; + constexpr size_t scale_plane_bytes = row_group * sizeof(ggml_fp16_t); + constexpr size_t field_bytes = row_group * 64; + constexpr size_t record_bytes = scale_plane_bytes + 4 * field_bytes; + constexpr size_t row_record_bytes = record_bytes / row_group; + + size_t expected_source_bytes = 0; + size_t expected_output_bytes = 0; + if (!checked_multiply(row_bytes, row_count, expected_source_bytes) || source.length != expected_source_bytes || + !checked_multiply(row_count, block_count, expected_output_bytes) || + !checked_multiply(expected_output_bytes, row_record_bytes, expected_output_bytes) || + source.materialized_length != expected_output_bytes) { + status.log("layout %s has inconsistent source/materialized lengths for K=%lld rows=%lld", source.layout.c_str(), + static_cast(source.input_size), static_cast(source.output_size)); + return status; + } + + output.resize(expected_output_bytes); + const auto * source_bytes = static_cast(source.host_data) + source.offset; + const size_t thread_count = + std::min(row_count, static_cast(std::max(1u, std::thread::hardware_concurrency()))); + std::atomic conversion_failed = false; + + auto convert_rows = [&](size_t row_begin, size_t row_end) { + std::array values = {}; + for (size_t row = row_begin; row < row_end; ++row) { + const size_t row_group_id = row / row_group; + const size_t lane = row % row_group; + const uint8_t * source_row = source_bytes + row * row_bytes; + for (size_t block = 0; block < block_count; ++block) { + dequantize_row_q5_K(reinterpret_cast(source_row) + block, values.data(), QK_K); + + float positive_max = 0.0f; + float negative_max = 0.0f; + for (float value : values) { + if (!std::isfinite(value)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + positive_max = std::max(positive_max, value); + negative_max = std::max(negative_max, -value); + } + float scale = std::max(positive_max, negative_max) / 127.0f; + if (scale == 0.0f) { + scale = 1.0f; + } + for (int iteration = 0; iteration < 2; ++iteration) { + double numerator = 0.0; + double denominator = 0.0; + for (float value : values) { + int quantized = 0; + if (!quantize_symmetric_value(value, scale, -127, 127, quantized)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + numerator += static_cast(value) * quantized; + denominator += static_cast(quantized) * quantized; + } + if (denominator > 0.0) { + scale = static_cast(numerator / denominator); + } + if (!std::isfinite(scale) || scale <= 0.0f) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + } + + const ggml_fp16_t encoded_scale = ggml_fp32_to_fp16(scale); + scale = ggml_fp16_to_fp32(encoded_scale); + const size_t record_base = (row_group_id * block_count + block) * record_bytes; + std::memcpy(output.data() + record_base + lane * sizeof(encoded_scale), &encoded_scale, + sizeof(encoded_scale)); + for (size_t chunk = 0; chunk < 4; ++chunk) { + auto * payload = reinterpret_cast(output.data() + record_base + scale_plane_bytes + + chunk * field_bytes + lane * 64); + for (size_t element = 0; element < 64; ++element) { + int quantized = 0; + if (!quantize_symmetric_value(values[chunk * 64 + element], scale, -127, 127, quantized)) { + conversion_failed.store(true, std::memory_order_relaxed); + return; + } + payload[element] = static_cast(quantized); + } + } + } + } + }; + + if (!run_worker_threads(thread_count, [&](size_t thread) { + convert_rows(row_count * thread / thread_count, row_count * (thread + 1) / thread_count); + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + return status; + } + if (conversion_failed.load(std::memory_order_relaxed)) { + output.clear(); + status.log("layout %s cannot represent non-finite symmetric weights", source.layout.c_str()); + } + return status; +} + +// Transpose native Q4_K records in 16-byte fields. No scale or code is +// reconstructed: applying the inverse permutation restores every source byte. +static Status materialize_q4_k_packed_k256_row64(const HostWeightSource & source, std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q4_K || source.input_size <= 0 || source.input_size % QK_K != 0 || + source.output_size <= 0 || source.output_size % 64 != 0) { + status.log("layout %s requires Q4_K, K divisible by %d, and rows divisible by 64", source.layout.c_str(), + QK_K); + return status; + } + + constexpr size_t row_group = 64; + constexpr size_t field_bytes = 16; + constexpr size_t field_count = sizeof(block_q4_K) / field_bytes; + static_assert(sizeof(block_q4_K) == 144); + const size_t blocks = static_cast(source.input_size / QK_K); + const size_t rows = static_cast(source.output_size); + size_t bytes = 0; + if (!checked_multiply(ggml_row_size(GGML_TYPE_Q4_K, source.input_size), rows, bytes) || + source.length != bytes || source.materialized_length != bytes) { + status.log("layout %s has inconsistent source/materialized lengths", source.layout.c_str()); + return status; + } + + output.resize(bytes); + const auto * input = static_cast(source.host_data) + source.offset; + const size_t groups = rows / row_group; + const size_t thread_count = std::min(groups, std::max(1u, std::thread::hardware_concurrency())); + if (!run_worker_threads(thread_count, [&](size_t thread) { + const size_t begin = groups * thread / thread_count; + const size_t end = groups * (thread + 1) / thread_count; + for (size_t group = begin; group < end; ++group) { + for (size_t block = 0; block < blocks; ++block) { + for (size_t field = 0; field < field_count; ++field) { + for (size_t lane = 0; lane < row_group; ++lane) { + const size_t src = (((group * row_group + lane) * blocks + block) * field_count + field) * + field_bytes; + const size_t dst = (((group * blocks + block) * field_count + field) * row_group + lane) * + field_bytes; + std::memcpy(output.data() + dst, input + src, field_bytes); + } + } + } + } + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + } + return status; +} + +static Status materialize_q6_k_i8_k32_row64(const HostWeightSource & source, std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q6_K) { + status.log("layout %s does not support GGML type %d", source.layout.c_str(), + static_cast(source.source_type)); + return status; + } + if (source.input_size <= 0 || source.input_size % QK_K != 0 || source.output_size <= 0 || + source.output_size % 64 != 0) { + status.log("layout %s requires K divisible by %d and rows divisible by 64, got K=%lld rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.input_size), + static_cast(source.output_size)); + return status; + } + + static_assert(sizeof(block_q6_K) == 210); + constexpr size_t row_group = 64; + constexpr size_t group_count = 8; + constexpr size_t group_elements = 32; + constexpr size_t d_plane_bytes = row_group * sizeof(ggml_fp16_t); + constexpr size_t scale_plane_bytes = group_count * row_group * 2; + constexpr size_t code_plane_bytes = group_count * row_group * group_elements; + constexpr size_t tile_block_bytes = d_plane_bytes + scale_plane_bytes + code_plane_bytes; + static_assert(tile_block_bytes == row_group * 274); + + const size_t block_count = static_cast(source.input_size / QK_K); + const size_t row_bytes = block_count * sizeof(block_q6_K); + size_t expected_source_bytes = 0; + size_t expected_output_bytes = 0; + if (!checked_multiply(row_bytes, static_cast(source.output_size), expected_source_bytes) || + !checked_multiply(static_cast(source.output_size), block_count, expected_output_bytes) || + !checked_multiply(expected_output_bytes, size_t{ 274 }, expected_output_bytes) || + source.length != expected_source_bytes || source.materialized_length != expected_output_bytes) { + status.log("layout %s source/materialized lengths %zu/%zu do not match expected %zu/%zu", source.layout.c_str(), + source.length, source.materialized_length, expected_source_bytes, expected_output_bytes); + return status; + } + + output.resize(expected_output_bytes); + const auto * source_bytes = static_cast(source.host_data) + source.offset; + const size_t output_group_count = static_cast(source.output_size) / row_group; + const size_t thread_count = std::min(output_group_count, std::max(1u, std::thread::hardware_concurrency())); + + auto convert_groups = [&](size_t group_begin, size_t group_end) { + for (size_t output_group = group_begin; output_group < group_end; ++output_group) { + for (size_t block = 0; block < block_count; ++block) { + const size_t tile_base = (output_group * block_count + block) * tile_block_bytes; + uint8_t * d_plane = output.data() + tile_base; + int8_t * scale_plane = reinterpret_cast(d_plane + d_plane_bytes); + int8_t * code_plane = scale_plane + scale_plane_bytes; + for (size_t lane = 0; lane < row_group; ++lane) { + const size_t row = output_group * row_group + lane; + const auto * source_block = reinterpret_cast(source_bytes + row * row_bytes + + block * sizeof(block_q6_K)); + std::memcpy(d_plane + lane * sizeof(ggml_fp16_t), &source_block->d, sizeof(ggml_fp16_t)); + for (size_t quant_group = 0; quant_group < group_count; ++quant_group) { + int8_t * scales = scale_plane + (quant_group * row_group + lane) * 2; + scales[0] = source_block->scales[quant_group * 2]; + scales[1] = source_block->scales[quant_group * 2 + 1]; + int8_t * codes = code_plane + (quant_group * row_group + lane) * group_elements; + const size_t half128 = quant_group / 4; + const size_t group_in_half = quant_group % 4; + for (size_t element = 0; element < group_elements; ++element) { + const size_t low_index = half128 * 64 + (group_in_half % 2) * 32 + element; + const size_t high_index = half128 * 32 + element; + const uint8_t low = + static_cast((source_block->ql[low_index] >> ((group_in_half / 2) * 4)) & 0x0F); + const uint8_t high = + static_cast((source_block->qh[high_index] >> (group_in_half * 2)) & 0x03); + codes[element] = static_cast((low | (high << 4)) - 32); + } + } + } + } + } + }; + + if (!run_worker_threads(thread_count, [&](size_t thread) { + const size_t group_begin = output_group_count * thread / thread_count; + const size_t group_end = output_group_count * (thread + 1) / thread_count; + convert_groups(group_begin, group_end); + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + } + return status; +} + +static Status materialize_q6_k_packed_k256_row64_scalerow(const HostWeightSource & source, + std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q6_K) { + status.log("layout %s does not support GGML type %d", source.layout.c_str(), + static_cast(source.source_type)); + return status; + } + if (source.input_size <= 0 || source.input_size % QK_K != 0 || source.output_size <= 0 || + source.output_size % 64 != 0) { + status.log("layout %s requires K divisible by %d and rows divisible by 64, got K=%lld rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.input_size), + static_cast(source.output_size)); + return status; + } + + static_assert(sizeof(block_q6_K) == 210); + constexpr size_t row_group = 64; + constexpr size_t half_count = 2; + constexpr size_t raw_field_count = 3; + constexpr size_t raw_field_bytes = 32; + constexpr size_t scale_group_count = 8; + constexpr size_t d_plane_bytes = row_group * sizeof(ggml_fp16_t); + constexpr size_t scale_plane_bytes = scale_group_count * row_group * half_count; + constexpr size_t raw_plane_bytes = half_count * raw_field_count * row_group * raw_field_bytes; + constexpr size_t tile_block_bytes = d_plane_bytes + scale_plane_bytes + raw_plane_bytes; + static_assert(tile_block_bytes == row_group * sizeof(block_q6_K)); + + const size_t block_count = static_cast(source.input_size / QK_K); + const size_t row_bytes = block_count * sizeof(block_q6_K); + size_t expected_source_bytes = 0; + if (!checked_multiply(row_bytes, static_cast(source.output_size), expected_source_bytes) || + source.length != expected_source_bytes || source.materialized_length != expected_source_bytes) { + status.log("layout %s source/materialized lengths %zu/%zu do not match expected %zu", source.layout.c_str(), + source.length, source.materialized_length, expected_source_bytes); + return status; + } + + output.resize(expected_source_bytes); + const auto * source_bytes = static_cast(source.host_data) + source.offset; + const size_t output_group_count = static_cast(source.output_size) / row_group; + const size_t thread_count = std::min(output_group_count, std::max(1u, std::thread::hardware_concurrency())); + + auto convert_groups = [&](size_t group_begin, size_t group_end) { + for (size_t output_group = group_begin; output_group < group_end; ++output_group) { + for (size_t block = 0; block < block_count; ++block) { + const size_t tile_base = (output_group * block_count + block) * tile_block_bytes; + uint8_t * d_plane = output.data() + tile_base; + int8_t * scale_plane = reinterpret_cast(d_plane + d_plane_bytes); + uint8_t * raw_plane = reinterpret_cast(scale_plane + scale_plane_bytes); + for (size_t lane = 0; lane < row_group; ++lane) { + const size_t row = output_group * row_group + lane; + const auto * source_block = reinterpret_cast(source_bytes + row * row_bytes + + block * sizeof(block_q6_K)); + std::memcpy(d_plane + lane * sizeof(ggml_fp16_t), &source_block->d, sizeof(ggml_fp16_t)); + for (size_t quant_group = 0; quant_group < scale_group_count; ++quant_group) { + for (size_t half = 0; half < half_count; ++half) { + const size_t destination_index = + (lane * half_count + half) * scale_group_count + quant_group; + scale_plane[destination_index] = source_block->scales[quant_group * half_count + half]; + } + } + for (size_t half = 0; half < half_count; ++half) { + uint8_t * ql0 = raw_plane + ((half * raw_field_count) * row_group + lane) * raw_field_bytes; + uint8_t * ql1 = ql0 + row_group * raw_field_bytes; + uint8_t * qh = ql1 + row_group * raw_field_bytes; + std::memcpy(ql0, source_block->ql + half * 64, raw_field_bytes); + std::memcpy(ql1, source_block->ql + half * 64 + raw_field_bytes, raw_field_bytes); + std::memcpy(qh, source_block->qh + half * raw_field_bytes, raw_field_bytes); + } + } + } + } + }; + + if (!run_worker_threads(thread_count, [&](size_t thread) { + const size_t group_begin = output_group_count * thread / thread_count; + const size_t group_end = output_group_count * (thread + 1) / thread_count; + convert_groups(group_begin, group_end); + })) { + output.clear(); + status.log("layout %s host materialization worker failed", source.layout.c_str()); + } + return status; +} + +static Status materialize_q6_k_symmetric_i2_packed_k256_row64_scalerow(const HostWeightSource & source, + std::vector & output) { + Status status; + if (source.source_type != GGML_TYPE_Q6_K || source.input_size <= 0 || source.input_size % QK_K != 0 || + source.output_size <= 0 || source.output_size % 64 != 0 || + static_cast(source.output_size) > std::numeric_limits::max() - 255) { + status.log( + "layout %s requires Q6_K, K divisible by %d, and rows divisible by 64, got type=%d K=%lld " + "rows=%lld", + source.layout.c_str(), QK_K, static_cast(source.source_type), + static_cast(source.input_size), static_cast(source.output_size)); + return status; + } + + const size_t logical_rows = static_cast(source.output_size); + const size_t row_group = ((logical_rows + 255) / 256) * 32; + const size_t physical_rows = (logical_rows + row_group - 1) / row_group * row_group; + const size_t block_count = static_cast(source.input_size / QK_K); + size_t symmetric_bytes = 0; + size_t packed_bytes = 0; + if (!checked_multiply(physical_rows, block_count, symmetric_bytes) || + !checked_multiply(symmetric_bytes, size_t{ 68 }, symmetric_bytes) || + !checked_multiply(logical_rows, block_count, packed_bytes) || + !checked_multiply(packed_bytes, sizeof(block_q6_K), packed_bytes) || + symmetric_bytes > std::numeric_limits::max() - packed_bytes || source.length != packed_bytes || + source.materialized_length != symmetric_bytes + packed_bytes) { + status.log("layout %s source/materialized lengths %zu/%zu are inconsistent with K=%lld rows=%lld", + source.layout.c_str(), source.length, source.materialized_length, + static_cast(source.input_size), static_cast(source.output_size)); + return status; + } + + HostWeightSource symmetric_source = source; + symmetric_source.layout = kSymmetricI2K32EightGroupsShared4Layout; + symmetric_source.materialized_length = symmetric_bytes; + status = materialize_symmetric_k32_eightgroups_shared4(symmetric_source, 2, output); + if (!status.success()) { + return status; + } + output.resize(source.materialized_length); + + HostWeightSource packed_source = source; + packed_source.layout = kQ6KPackedK256Row64ScaleRowLayout; + packed_source.materialized_length = packed_bytes; + std::vector packed; + status = materialize_q6_k_packed_k256_row64_scalerow(packed_source, packed); + if (!status.success()) { + output.clear(); + return status; + } + std::memcpy(output.data() + symmetric_bytes, packed.data(), packed.size()); + return status; +} + +static Status materialize_weight(const HostWeightSource & source, + const void *& upload_data, + size_t & upload_size, + std::vector & transformed) { + Status status; + if (source.layout == kNativeWeightLayout) { + upload_data = static_cast(source.host_data) + source.offset; + upload_size = source.length; + return status; + } + if (source.layout == kSymmetricI4K32EightGroupsShared4Layout || + source.layout == kSymmetricI4K32EightGroupsShared4MultistartLayout) { + status = materialize_symmetric_i4_k32_eightgroups_shared4(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kSymmetricI4K32Row64Layout) { + status = materialize_symmetric_i4_k32_row64(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kSymmetricI4K64Row64Layout) { + status = materialize_symmetric_i4_k64_row64(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kQ5KSymmetricI5K32Layout) { + status = materialize_symmetric_i5_k32(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kQ5KSymmetricI8K256Row64Layout) { + status = materialize_symmetric_i8_k256_row64(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kQ4KPackedK256Row64Layout) { + status = materialize_q4_k_packed_k256_row64(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kQ6KI8K32Row64Layout) { + status = materialize_q6_k_i8_k32_row64(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kQ6KPackedK256Row64ScaleRowLayout) { + status = materialize_q6_k_packed_k256_row64_scalerow(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + if (source.layout == kQ6KSymmetricI2PackedK256Row64ScaleRowLayout) { + status = materialize_q6_k_symmetric_i2_packed_k256_row64_scalerow(source, transformed); + if (status.success()) { + upload_data = transformed.data(); + upload_size = transformed.size(); + } + return status; + } + status.log("unknown host weight layout %s", source.layout.c_str()); + return status; +} static Status allocate_device_buffer(hrx_device_t device, size_t size, hrx_buffer_t & buffer) { Status status; @@ -36,6 +1227,44 @@ static Status allocate_device_buffer(hrx_device_t device, size_t size, hrx_buffe } // namespace +Status allocate_mapped_host_staging_buffer(hrx_device_t device, size_t size, hrx_buffer_t & buffer, void *& host_data) { + Status status; + buffer = nullptr; + host_data = nullptr; + if (device == nullptr) { + status.log("missing HRX device for host staging allocation"); + return status; + } + if (size == 0) { + status.log("cannot allocate an empty host staging buffer"); + return status; + } + hrx_buffer_params_t params = { + HRX_MEMORY_TYPE_HOST_LOCAL | HRX_MEMORY_TYPE_DEVICE_VISIBLE, + HRX_MEMORY_ACCESS_ALL, + HRX_BUFFER_USAGE_DEFAULT | HRX_BUFFER_USAGE_MAPPING_SCOPED | HRX_BUFFER_USAGE_MAPPING_PERSISTENT, + 0, + }; + if (ErrorResult error = + take_status(hrx_allocator_allocate_buffer(hrx_device_allocator(device), params, size, &buffer))) { + status.log("allocate mapped HRX host staging buffer: %s", error->c_str()); + return status; + } + if (ErrorResult error = take_status(hrx_buffer_map(buffer, HRX_MAP_READ | HRX_MAP_WRITE, 0, size, &host_data))) { + status.log("map HRX host staging buffer: %s", error->c_str()); + hrx_buffer_release(buffer); + buffer = nullptr; + } + return status; +} + +HostTransferManager::~HostTransferManager() { + std::lock_guard lock(mutex_); + if (mapped_download_buffer_ != nullptr) { + hrx_buffer_release(mapped_download_buffer_); + } +} + Status HostTransferManager::upload_synchronous(hrx_stream_t stream, const void * host_source, hrx_buffer_t destination, @@ -82,6 +1311,12 @@ Status HostTransferManager::upload_async(hrx_stream_t stream, return status; } + // Inline stream updates are intended for small payloads. Use bulk H2D for + // large host-staged tensors to avoid hundreds of update commands. + if (size >= kLargeHostUploadBytes) { + return upload_synchronous(stream, host_source, destination, offset, size); + } + const uint8_t * host_bytes = static_cast(host_source); size_t uploaded = 0; while (uploaded < size) { @@ -133,6 +1368,55 @@ Status HostTransferManager::download_synchronous(hrx_stream_t stream, return status; } +Status HostTransferManager::download_to_mapped_synchronous(hrx_device_t device, + hrx_stream_t stream, + hrx_buffer_t source, + size_t offset, + size_t size, + HostMappedDownloadLease & lease) { + Status status; + lease = {}; + if (device == nullptr || stream == nullptr || source == nullptr || size == 0) { + status.log("invalid mapped HRX host download"); + return status; + } + lease.lock = std::unique_lock(mutex_); + if (mapped_download_capacity_ < size) { + size_t capacity = std::max(mapped_download_capacity_, 1); + while (capacity < size && capacity <= std::numeric_limits::max() / 2) { + capacity *= 2; + } + if (capacity < size) { + capacity = size; + } + hrx_buffer_t new_buffer = nullptr; + void * new_data = nullptr; + status = allocate_mapped_host_staging_buffer(device, capacity, new_buffer, new_data); + if (!status.success()) { + return status; + } + if (mapped_download_buffer_ != nullptr) { + hrx_buffer_release(mapped_download_buffer_); + } + mapped_download_buffer_ = new_buffer; + mapped_download_data_ = new_data; + mapped_download_capacity_ = capacity; + } + if (ErrorResult error = + take_status(hrx_stream_copy_buffer(stream, source, offset, mapped_download_buffer_, 0, size))) { + status.log("mapped host weight download failed: %s", error->c_str()); + return status; + } + if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { + status.log("synchronize mapped host weight download failed: %s", error->c_str()); + return status; + } + ++stats_.downloads; + stats_.download_bytes += size; + lease.host_data = mapped_download_data_; + return status; +} + HostTransferStats HostTransferManager::stats() const { std::lock_guard lock(mutex_); return stats_; @@ -189,6 +1473,13 @@ size_t HostWeightCache::SourceKeyHash::operator()(const SourceKey & key) const { mix(static_cast(key.capacity)); mix(static_cast(key.offset)); mix(static_cast(key.length)); + mix(static_cast(key.materialized_length)); + for (unsigned char byte : key.layout) { + mix(byte); + } + mix(static_cast(key.source_type)); + mix(static_cast(key.input_size)); + mix(static_cast(key.output_size)); return static_cast(hash); } @@ -201,7 +1492,9 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device result.status.log("host weight residency requires an HRX stream"); return result; } - if (source.host_data == nullptr || source.identity == 0 || source.generation == 0 || source.length == 0 || + const bool has_host_source = source.host_data != nullptr; + const bool has_device_source = source.device_buffer != nullptr; + if (has_host_source == has_device_source || source.identity == 0 || source.generation == 0 || source.length == 0 || source.offset > source.capacity || source.length > source.capacity - source.offset) { result.status.log("invalid host weight source"); return result; @@ -211,32 +1504,79 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device return result; } - const SourceKey key{ source.identity, source.generation, source.capacity, source.offset, source.length }; + const size_t materialized_length = + source.layout == kNativeWeightLayout ? source.length : source.materialized_length; + if (materialized_length == 0) { + result.status.log("host weight source has an empty materialized layout"); + return result; + } + + const SourceKey key{ source.identity, source.generation, source.capacity, source.offset, + source.length, materialized_length, source.layout, source.source_type, + source.input_size, source.output_size }; { std::lock_guard lock(mutex_); const auto found = entries_.find(key); if (found != entries_.end()) { - if (found->second->layout != source.layout) { - ++stats_.layout_conflicts; - result.status.log("host weight source already has resident layout %s, cannot also materialize %s", - found->second->layout.c_str(), source.layout.c_str()); - return result; - } ++stats_.hits; result.lease = HostWeightLease(found->second); return result; } } + HostWeightSource materialization_source = source; + HostMappedDownloadLease mapped_canonical; + std::vector canonical; + if (has_device_source) { + if (source.layout == kNativeWeightLayout) { + result.status.log("native device weights do not require host materialization"); + return result; + } + if (!mapped_download_staging_enabled()) { + canonical.resize(source.length); + } + void * canonical_data = canonical.data(); + if (mapped_download_staging_enabled()) { + result.status = transfers.download_to_mapped_synchronous( + device, stream, source.device_buffer, source.offset, source.length, mapped_canonical); + canonical_data = mapped_canonical.host_data; + } else { + result.status = transfers.download_synchronous(stream, source.device_buffer, source.offset, canonical_data, + source.length); + } + if (!result.status.success()) { + return result; + } + materialization_source.host_data = canonical_data; + materialization_source.device_buffer = nullptr; + materialization_source.capacity = source.length; + materialization_source.offset = 0; + } + + const void * upload_data = nullptr; + size_t upload_size = 0; + std::vector transformed; + result.status = materialize_weight(materialization_source, upload_data, upload_size, transformed); + if (mapped_canonical.lock.owns_lock()) { + mapped_canonical.lock.unlock(); + } + if (!result.status.success()) { + return result; + } + if (upload_size != materialized_length) { + result.status.log("host weight layout %s produced %zu bytes, expected %zu", source.layout.c_str(), upload_size, + materialized_length); + return result; + } + auto entry = std::make_shared(); - entry->length = source.length; + entry->length = materialized_length; entry->layout = source.layout; - result.status = allocate_device_buffer(device, source.length, entry->buffer); + result.status = allocate_device_buffer(device, materialized_length, entry->buffer); if (!result.status.success()) { return result; } - result.status = transfers.upload_synchronous( - stream, static_cast(source.host_data) + source.offset, entry->buffer, 0, source.length); + result.status = transfers.upload_synchronous(stream, upload_data, entry->buffer, 0, upload_size); if (!result.status.success()) { return result; } @@ -251,7 +1591,7 @@ HostWeightAcquireResult HostWeightCache::acquire(hrx_device_t device } ++stats_.misses; stats_.allocation_count = entries_.size(); - stats_.resident_bytes += source.length; + stats_.resident_bytes += materialized_length; } result.lease = HostWeightLease(std::move(entry)); return result; @@ -283,6 +1623,7 @@ HostStagingBuffer & HostStagingBuffer::operator=(HostStagingBuffer && other) noe clear(); buffer = other.buffer; host_data = other.host_data; + source_host_buffer = std::move(other.source_host_buffer); value = other.value; length = other.length; upload = other.upload; @@ -302,6 +1643,7 @@ void HostStagingBuffer::clear() { buffer = nullptr; } host_data = nullptr; + source_host_buffer = HostBufferRef{}; value = -1; length = 0; upload = false; diff --git a/ggml/src/ggml-hrx/runtime/host-memory.h b/ggml/src/ggml-hrx/runtime/host-memory.h index ee4dc45efec4..a125598f20bf 100644 --- a/ggml/src/ggml-hrx/runtime/host-memory.h +++ b/ggml/src/ggml-hrx/runtime/host-memory.h @@ -1,5 +1,8 @@ #pragma once +#include "dispatch/dispatch.h" +#include "ggml.h" +#include "runtime/host-buffer-registry.h" #include "status.h" #include @@ -22,17 +25,36 @@ struct HostTransferStats { size_t download_bytes = 0; }; +struct HostMappedDownloadLease { + std::unique_lock lock; + void * host_data = nullptr; +}; + class HostTransferManager { public: - Status upload_synchronous( - hrx_stream_t stream, const void * host_source, hrx_buffer_t destination, size_t offset, size_t size); + ~HostTransferManager(); + + Status upload_synchronous(hrx_stream_t stream, + const void * host_source, + hrx_buffer_t destination, + size_t offset, + size_t size); Status upload_async(hrx_stream_t stream, const void * host_source, hrx_buffer_t destination, size_t offset, size_t size); - Status download_synchronous( - hrx_stream_t stream, hrx_buffer_t source, size_t offset, void * host_destination, size_t size); + Status download_synchronous(hrx_stream_t stream, + hrx_buffer_t source, + size_t offset, + void * host_destination, + size_t size); + Status download_to_mapped_synchronous(hrx_device_t device, + hrx_stream_t stream, + hrx_buffer_t source, + size_t offset, + size_t size, + HostMappedDownloadLease & lease); HostTransferStats stats() const; void clear(); @@ -40,22 +62,29 @@ class HostTransferManager { private: mutable std::mutex mutex_; HostTransferStats stats_; + hrx_buffer_t mapped_download_buffer_ = nullptr; + void * mapped_download_data_ = nullptr; + size_t mapped_download_capacity_ = 0; }; struct HostWeightSource { const void * host_data = nullptr; + hrx_buffer_t device_buffer = nullptr; uint64_t identity = 0; uint64_t generation = 0; size_t capacity = 0; size_t offset = 0; size_t length = 0; - std::string layout = "ggml-native"; + size_t materialized_length = 0; + std::string layout = kNativeWeightLayout; + ggml_type source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; }; struct HostWeightCacheStats { uint64_t hits = 0; uint64_t misses = 0; - uint64_t layout_conflicts = 0; size_t allocation_count = 0; size_t resident_bytes = 0; }; @@ -106,10 +135,18 @@ class HostWeightCache { size_t capacity = 0; size_t offset = 0; size_t length = 0; + size_t materialized_length = 0; + std::string layout; + ggml_type source_type = GGML_TYPE_COUNT; + int64_t input_size = 0; + int64_t output_size = 0; bool operator==(const SourceKey & other) const { return identity == other.identity && generation == other.generation && capacity == other.capacity && - offset == other.offset && length == other.length; + offset == other.offset && length == other.length && + materialized_length == other.materialized_length && layout == other.layout && + source_type == other.source_type && input_size == other.input_size && + output_size == other.output_size; } }; @@ -134,6 +171,7 @@ struct HostStagingBuffer { hrx_buffer_t buffer = nullptr; void * host_data = nullptr; + HostBufferRef source_host_buffer; int32_t value = -1; size_t length = 0; bool upload = false; @@ -143,5 +181,9 @@ struct HostStagingBuffer { }; Status allocate_host_staging_buffer(hrx_device_t device, size_t size, HostStagingBuffer & staging); +Status allocate_mapped_host_staging_buffer(hrx_device_t device, + size_t size, + hrx_buffer_t & buffer, + void *& host_data); } // namespace ggml::hrx 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 8d107144aa9b..d79a1fbf5cfe 100644 --- a/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp +++ b/ggml/src/ggml-hrx/runtime/loom-kernel-jit.cpp @@ -26,6 +26,8 @@ class LoomAmdgpuJit { ggml_hrx_loom_jit_amdgpu_options options = {}; options.processor = target; options.identifier = target; + options.sanitizer = std::getenv("GGML_HRX_LOOM_SANITIZER"); + options.sanitizer_reporting = std::getenv("GGML_HRX_LOOM_SANITIZER_REPORTING"); if (ErrorResult error = take_status(ggml_hrx_loom_jit_amdgpu_create(&options, &jit_))) { error_message = "create Loom JIT: " + *error; GGML_LOG_ERROR("%s: %s\n", __func__, error_message.c_str()); @@ -48,11 +50,20 @@ class LoomAmdgpuJit { }; static size_t default_worker_count() { + const char * configured = std::getenv("GGML_HRX_JIT_WORKERS"); + if (configured != nullptr && configured[0] != '\0') { + char * end = nullptr; + const unsigned long count = std::strtoul(configured, &end, 10); + if (end != configured && *end == '\0' && count >= 1 && count <= 64) { + return static_cast(count); + } + GGML_LOG_WARN("%s: ignoring invalid GGML_HRX_JIT_WORKERS=%s\n", __func__, configured); + } const unsigned int hardware_threads = std::thread::hardware_concurrency(); if (hardware_threads == 0) { return 1; } - return std::min(hardware_threads, 4); + return std::min(hardware_threads, 8); } static bool compile_kernel(ggml_hrx_loom_jit_amdgpu * jit, @@ -79,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/runtime/prepared-command-program-cache.cpp b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp index ba8c7db983fe..ad81b1a23afc 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.cpp @@ -25,14 +25,6 @@ static uint64_t hash_text(const char * text) { return hash; } -static uint64_t hash_string(const std::string & text) { - uint64_t hash = UINT64_C(1469598103934665603); - for (const char c : text) { - mix_hash(hash, static_cast(c)); - } - return hash; -} - static void apply_graph_replay_result(PreparedCommandProgramCacheExecutionResult & result, const RecordedCommandGraphExecutionResult & replay) { result.graph_replay_event = replay.event; @@ -50,6 +42,14 @@ static bool graph_replay_should_fallback(HrxGraphReplayEvent event) { } // namespace +uint64_t command_program_shape_hash(const std::string & command_shape) { + uint64_t hash = UINT64_C(1469598103934665603); + for (const char c : command_shape) { + mix_hash(hash, static_cast(c)); + } + return hash; +} + size_t PreparedCommandProgramCache::KeyHash::operator()(const Key & key) const { uint64_t hash = UINT64_C(1469598103934665603); mix_hash(hash, key.graph_uid); @@ -61,12 +61,12 @@ size_t PreparedCommandProgramCache::KeyHash::operator()(const Key & key) const { PreparedCommandProgramCache::Key PreparedCommandProgramCache::cache_key(uint64_t graph_uid, const CommandProgramExecutionContext & context, - const std::string & command_shape, + uint64_t command_shape_hash, const CommandProgramBindings & bindings) const { return { graph_uid, hash_text(context.target), - hash_string(command_shape), + command_shape_hash, command_program_bindings_hash(bindings).value, }; } @@ -85,6 +85,15 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ const std::string & command_shape, const CommandProgram & commands, const CommandProgramBindings & bindings) { + return execute_with_result(context, graph_uid, command_program_shape_hash(command_shape), commands, bindings); +} + +PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_with_result( + const CommandProgramExecutionContext & context, + uint64_t graph_uid, + uint64_t command_shape_hash, + const CommandProgram & commands, + const CommandProgramBindings & bindings) { PreparedCommandProgramCacheExecutionResult result; if (graph_uid == 0 || !commands.valid() || !bindings.valid()) { result.graph_replay_event = HrxGraphReplayEvent::Ineligible; @@ -101,7 +110,7 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ return result; } - const Key key = cache_key(graph_uid, context, command_shape, bindings); + const Key key = cache_key(graph_uid, context, command_shape_hash, bindings); std::shared_ptr entry; bool created_entry = false; { @@ -119,6 +128,15 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ std::lock_guard entry_lock(entry->mutex); if (entry->has_program && entry->program.valid()) { record_hit(); + if (debug_serial_command_execution_enabled()) { + result.graph_replay_event = HrxGraphReplayEvent::Disabled; + result.graph_replay_ineligible_reason = "debug_serial_execution"; + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + if (!result.success) { + result.status.log("execute cached HRX command program failed"); + } + return result; + } const RecordedCommandGraphExecutionResult replay = bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->recorded); apply_graph_replay_result(result, replay); @@ -156,6 +174,16 @@ PreparedCommandProgramCacheExecutionResult PreparedCommandProgramCache::execute_ entry->has_program = true; record_build(); + if (debug_serial_command_execution_enabled()) { + result.graph_replay_event = HrxGraphReplayEvent::Disabled; + result.graph_replay_ineligible_reason = "debug_serial_execution"; + result.success = bind_and_execute_prepared_command_program(context, commands, bindings, entry->program); + if (!result.success) { + result.status.log("execute prepared HRX command program failed"); + } + return result; + } + const RecordedCommandGraphExecutionResult replay = bind_and_launch_recorded_command_graph(context, commands, bindings, entry->program, entry->recorded); apply_graph_replay_result(result, replay); diff --git a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h index 68b61b0ce8f2..069ffd624480 100644 --- a/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/prepared-command-program-cache.h @@ -14,6 +14,8 @@ namespace ggml::hrx { +uint64_t command_program_shape_hash(const std::string & command_shape); + struct PreparedCommandProgramCacheStats { uint64_t builds = 0; uint64_t hits = 0; @@ -45,6 +47,12 @@ class PreparedCommandProgramCache { const CommandProgram & commands, const CommandProgramBindings & bindings); + PreparedCommandProgramCacheExecutionResult execute_with_result(const CommandProgramExecutionContext & context, + uint64_t graph_uid, + uint64_t command_shape_hash, + const CommandProgram & commands, + const CommandProgramBindings & bindings); + PreparedCommandProgramCacheStats stats() const; void clear(); @@ -68,7 +76,7 @@ class PreparedCommandProgramCache { Key cache_key(uint64_t graph_uid, const CommandProgramExecutionContext & context, - const std::string & command_shape, + uint64_t command_shape_hash, const CommandProgramBindings & bindings) const; struct Entry { diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.cpp b/ggml/src/ggml-hrx/runtime/transient-arena.cpp index a950f4ab8c76..a12f89e0ab38 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.cpp +++ b/ggml/src/ggml-hrx/runtime/transient-arena.cpp @@ -3,6 +3,7 @@ #include "hrx-interop-utils.h" #include "hrx_runtime.h" +#include #include namespace ggml::hrx { @@ -39,6 +40,10 @@ void TransientArena::clear() { hrx_buffer_release(buffer_); buffer_ = nullptr; } + for (hrx_buffer_t buffer : diagnostic_retired_buffers_) { + hrx_buffer_release(buffer); + } + diagnostic_retired_buffers_.clear(); allocation_capacity_ = 0; allocation_id_ = kInvalidTransientArenaAllocationId; } @@ -62,7 +67,8 @@ TransientArena::AllocationLease TransientArena::acquire_allocation_lease() { Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t stream, size_t required_size) { Status status; - if (required_size == 0 || allocation_capacity_ >= required_size) { + const bool diagnostic_fresh = std::getenv("GGML_HRX_DIAGNOSTIC_FRESH_TRANSIENT_ARENA") != nullptr; + if (required_size == 0 || (!diagnostic_fresh && allocation_capacity_ >= required_size)) { return status; } if (device == nullptr) { @@ -74,11 +80,15 @@ Status TransientArena::ensure_capacity_locked(hrx_device_t device, hrx_stream_t return status; } if (buffer_ != nullptr) { + if (diagnostic_fresh) { + diagnostic_retired_buffers_.push_back(buffer_); + } else { if (ErrorResult error = take_status(hrx_stream_synchronize(stream))) { status.log("synchronize before growing transient arena: %s", error->c_str()); return status; } hrx_buffer_release(buffer_); + } buffer_ = nullptr; allocation_capacity_ = 0; allocation_id_ = kInvalidTransientArenaAllocationId; diff --git a/ggml/src/ggml-hrx/runtime/transient-arena.h b/ggml/src/ggml-hrx/runtime/transient-arena.h index a529bf938f10..a50639962c7c 100644 --- a/ggml/src/ggml-hrx/runtime/transient-arena.h +++ b/ggml/src/ggml-hrx/runtime/transient-arena.h @@ -6,6 +6,7 @@ #include #include #include +#include typedef struct hrx_device_s * hrx_device_t; typedef struct hrx_stream_s * hrx_stream_t; @@ -58,6 +59,7 @@ class TransientArena { mutable std::mutex mutex_; hrx_buffer_t buffer_ = nullptr; + std::vector diagnostic_retired_buffers_; size_t allocation_capacity_ = 0; uint64_t allocation_id_ = kInvalidTransientArenaAllocationId; uint64_t next_allocation_id_ = 1; diff --git a/ggml/src/ggml-hrx/tools/benchmarks/analyze-model-fusion-adjacency.py b/ggml/src/ggml-hrx/tools/benchmarks/analyze-model-fusion-adjacency.py new file mode 100644 index 000000000000..00fa5edbe3b0 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/analyze-model-fusion-adjacency.py @@ -0,0 +1,750 @@ +#!/usr/bin/env python3 +# +# Analyze HRX command-program adjacency for model-scoped Loom benchmarks. + +from __future__ import annotations + +import argparse +import hashlib +import json +from collections import Counter, defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +TRANSIENT_ORIGIN = "Transient" +READ_ACCESS = "Read" +WRITE_ACCESSES = {"Write", "ReadWrite"} + + +@dataclass(frozen=True) +class BindingKey: + value: int + offset: int + length: int + + +@dataclass(frozen=True) +class Producer: + ordinal: int + kernel: str + binding: str + key: BindingKey + + +@dataclass(frozen=True) +class Consumer: + ordinal: int + kernel: str + binding: str + + +def fail(message: str) -> None: + raise SystemExit(message) + + +def load_json(path: Path) -> Any: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def write_json(path: Path, data: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, indent=2, sort_keys=True) + f.write("\n") + + +def kernel_symbol(kernel: str) -> str: + return kernel.split(":")[-1] + + +def command_shape_key(command: dict[str, Any]) -> str: + shape_data = { + "kernel": command.get("kernel"), + "integer_parameters": command.get("integer_parameters", {}), + "compile_parameters": command.get("compile_parameters", {}), + "binding_lengths": [binding.get("length") for binding in command.get("bindings", [])], + } + encoded = json.dumps(shape_data, sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest()[:16] + + +def load_programs(dump_dir: Path) -> list[dict[str, Any]]: + paths = sorted(dump_dir.glob("program-*/program.json")) + if not paths: + fail(f"no program.json files found under {dump_dir}") + programs = [] + for path in paths: + program = load_json(path) + program["path"] = str(path) + programs.append(program) + return programs + + +def load_commands(dump_dir: Path) -> list[dict[str, Any]]: + commands = [] + for program in load_programs(dump_dir): + for command in program.get("commands", []): + command = dict(command) + command["program"] = { + "directory": Path(program["path"]).parent.name, + "dump_id": program.get("dump_id"), + "shape_hash": program.get("shape_hash"), + "target": program.get("target"), + } + commands.append(command) + return commands + + +def shape_key_order(commands: list[dict[str, Any]]) -> list[str]: + keys = [] + seen = set() + for command in commands: + key = command_shape_key(command) + if key not in seen: + seen.add(key) + keys.append(key) + return keys + + +def shape_counts(commands: list[dict[str, Any]]) -> Counter[str]: + return Counter(command_shape_key(command) for command in commands) + + +def shape_metrics_from_summary(summary: dict[str, Any]) -> dict[str, dict[str, Any]]: + by_benchmark = {} + for kernel in summary.get("kernels", []): + for shape in kernel.get("shapes", []): + benchmark = shape.get("benchmark") + if benchmark: + by_benchmark[benchmark] = shape + return by_benchmark + + +def map_shape_keys_to_dispatches(commands: list[dict[str, Any]], + manifest: dict[str, Any], + summary: dict[str, Any]) -> tuple[dict[str, dict[str, Any]], list[str]]: + warnings = [] + keys = shape_key_order(commands) + dispatches = manifest.get("dispatches", []) + if len(keys) != len(dispatches): + warnings.append(f"shape count mismatch: dump has {len(keys)} unique shapes, manifest has {len(dispatches)} dispatches") + + summary_by_benchmark = shape_metrics_from_summary(summary) + shape_info = {} + for index, key in enumerate(keys[:len(dispatches)]): + dispatch = dispatches[index] + command = next(command for command in commands if command_shape_key(command) == key) + if command.get("kernel") != dispatch.get("kernel"): + warnings.append( + f"shape {index} kernel mismatch: dump has {command.get('kernel')}, manifest has {dispatch.get('kernel')}" + ) + benchmark = dispatch.get("benchmark") + metric = summary_by_benchmark.get(benchmark, {}) + shape_info[key] = { + "benchmark": benchmark, + "kernel": dispatch.get("kernel"), + "count": dispatch.get("count", 0), + "metric_ns": metric.get("metric_ns"), + "weighted_ns": metric.get("weighted_ns"), + "state": metric.get("state", "missing_summary"), + "error": metric.get("error"), + } + return shape_info, warnings + + +def command_metric_ns(command: dict[str, Any], shape_info: dict[str, dict[str, Any]]) -> float: + info = shape_info.get(command_shape_key(command), {}) + metric = info.get("metric_ns") + if metric is None: + return 0.0 + return float(metric) + + +def binding_key(binding: dict[str, Any]) -> BindingKey: + return BindingKey(int(binding["value"]), int(binding.get("offset", 0)), int(binding.get("length", 0))) + + +def read_bindings(command: dict[str, Any]) -> list[dict[str, Any]]: + return [ + binding + for binding in command.get("bindings", []) + if binding.get("origin") == TRANSIENT_ORIGIN and binding.get("access") == READ_ACCESS + ] + + +def write_bindings(command: dict[str, Any]) -> list[dict[str, Any]]: + return [ + binding + for binding in command.get("bindings", []) + if binding.get("origin") == TRANSIENT_ORIGIN and binding.get("access") in WRITE_ACCESSES + ] + + +def exact_transient_edges(commands: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[Producer, list[Consumer]]]: + active: dict[BindingKey, Producer] = {} + edges = [] + fanout: dict[Producer, list[Consumer]] = defaultdict(list) + for command in sorted(commands, key=lambda item: int(item["ordinal"])): + consumer_ordinal = int(command["ordinal"]) + for binding in read_bindings(command): + key = binding_key(binding) + producer = active.get(key) + if producer is None: + continue + consumer = Consumer(consumer_ordinal, command["kernel"], str(binding.get("name", ""))) + edges.append( + { + "producer_ordinal": producer.ordinal, + "producer_kernel": producer.kernel, + "producer_binding": producer.binding, + "consumer_ordinal": consumer.ordinal, + "consumer_kernel": consumer.kernel, + "consumer_binding": consumer.binding, + "value": key.value, + "offset": key.offset, + "length": key.length, + } + ) + fanout[producer].append(consumer) + for binding in write_bindings(command): + key = binding_key(binding) + active[key] = Producer(consumer_ordinal, command["kernel"], str(binding.get("name", "")), key) + return edges, fanout + + +def summarize_edges(edges: list[dict[str, Any]]) -> list[dict[str, Any]]: + counts: Counter[tuple[str, str]] = Counter() + examples: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + for edge in edges: + key = (edge["producer_kernel"], edge["consumer_kernel"]) + counts[key] += 1 + if len(examples[key]) < 5: + examples[key].append(edge) + return [ + { + "producer_kernel": producer, + "consumer_kernel": consumer, + "edge_count": count, + "examples": examples[(producer, consumer)], + } + for (producer, consumer), count in counts.most_common() + ] + + +def summarize_fanout(fanout: dict[Producer, list[Consumer]]) -> list[dict[str, Any]]: + counts: Counter[tuple[str, tuple[str, ...]]] = Counter() + examples: dict[tuple[str, tuple[str, ...]], list[dict[str, Any]]] = defaultdict(list) + for producer, consumers in fanout.items(): + consumer_kernels = tuple(sorted(consumer.kernel for consumer in consumers)) + key = (producer.kernel, consumer_kernels) + counts[key] += 1 + if len(examples[key]) < 5: + examples[key].append( + { + "producer_ordinal": producer.ordinal, + "producer_binding": producer.binding, + "consumers": [ + { + "ordinal": consumer.ordinal, + "kernel": consumer.kernel, + "binding": consumer.binding, + } + for consumer in consumers + ], + } + ) + rows = [] + for (producer_kernel, consumer_kernels), count in counts.most_common(): + rows.append( + { + "producer_kernel": producer_kernel, + "consumer_kernels": list(consumer_kernels), + "producer_count": count, + "examples": examples[(producer_kernel, consumer_kernels)], + } + ) + return rows + + +def producers_by_exact_key(commands: list[dict[str, Any]]) -> dict[tuple[int, BindingKey], Producer]: + active: dict[BindingKey, Producer] = {} + result = {} + for command in sorted(commands, key=lambda item: int(item["ordinal"])): + ordinal = int(command["ordinal"]) + for binding in read_bindings(command): + key = binding_key(binding) + producer = active.get(key) + if producer is not None: + result[(ordinal, key)] = producer + for binding in write_bindings(command): + key = binding_key(binding) + active[key] = Producer(ordinal, command["kernel"], str(binding.get("name", "")), key) + return result + + +def producer_for_read(command: dict[str, Any], + binding_name: str, + read_producers: dict[tuple[int, BindingKey], Producer]) -> Producer | None: + ordinal = int(command["ordinal"]) + for binding in read_bindings(command): + if binding.get("name") == binding_name: + return read_producers.get((ordinal, binding_key(binding))) + return None + + +def output_consumers(command: dict[str, Any], fanout: dict[Producer, list[Consumer]], binding_name: str = "output") -> list[Consumer]: + ordinal = int(command["ordinal"]) + for binding in write_bindings(command): + if binding.get("name") != binding_name: + continue + producer = Producer(ordinal, command["kernel"], str(binding.get("name", "")), binding_key(binding)) + return fanout.get(producer, []) + return [] + + +def command_by_ordinal(commands: list[dict[str, Any]]) -> dict[int, dict[str, Any]]: + return {int(command["ordinal"]): command for command in commands} + + +def compile_parameters(command: dict[str, Any]) -> dict[str, Any]: + return command.get("compile_parameters", {}) + + +def integer_parameters(command: dict[str, Any]) -> dict[str, Any]: + return command.get("integer_parameters", {}) + + +def op_name(command: dict[str, Any]) -> str: + if command["kernel"] != "loom_libs:ggml_binary_f32": + return "" + op = str(compile_parameters(command).get("ggml.binary_f32.op", "")) + names = { + "0": "add", + "1": "sub", + "2": "mul", + "3": "div", + "4": "swiglu", + "5": "geglu", + "6": "reglu", + "7": "geglu_erf", + "8": "geglu_quick", + } + return names.get(op, f"op_{op}") + + +def add_candidate(candidates: list[dict[str, Any]], + name: str, + kind: str, + occurrences: list[dict[str, Any]], + current_chain_ns: float, + removable_ns: float, + notes: list[str]) -> None: + candidates.append( + { + "name": name, + "kind": kind, + "occurrence_count": len(occurrences), + "current_chain_ns": current_chain_ns, + "removable_standalone_ns": removable_ns, + "examples": occurrences[:8], + "notes": notes, + } + ) + + +def classify_candidates(commands: list[dict[str, Any]], + shape_info: dict[str, dict[str, Any]], + fanout: dict[Producer, list[Consumer]]) -> list[dict[str, Any]]: + by_ordinal = command_by_ordinal(commands) + read_producers = producers_by_exact_key(commands) + candidates: list[dict[str, Any]] = [] + + attention_occurrences = [] + attention_chain_ns = 0.0 + attention_removable_ns = 0.0 + for command in commands: + if command["kernel"] != "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8": + continue + ordinal = int(command["ordinal"]) + rope = by_ordinal.get(ordinal - 3) + rope_set_rows = by_ordinal.get(ordinal - 2) + set_rows = by_ordinal.get(ordinal - 1) + q_projection = producer_for_read(rope, "input", read_producers) if rope else None + k_projection = producer_for_read(rope_set_rows, "input", read_producers) if rope_set_rows else None + v_projection = producer_for_read(set_rows, "rows", read_producers) if set_rows else None + matched = ( + rope is not None + and rope_set_rows is not None + and set_rows is not None + and rope["kernel"] == "loom_libs:ggml_rope_f32" + and rope_set_rows["kernel"] == "loom_libs:ggml_rope_set_rows_f32" + and set_rows["kernel"] == "loom_libs:ggml_set_rows" + and q_projection is not None + and k_projection is not None + and v_projection is not None + and q_projection.kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" + and k_projection.kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" + and v_projection.kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" + ) + if not matched: + continue + occurrence_ns = sum(command_metric_ns(item, shape_info) for item in (rope, rope_set_rows, set_rows)) + attention_chain_ns += occurrence_ns + command_metric_ns(command, shape_info) + attention_removable_ns += occurrence_ns + attention_occurrences.append( + { + "attention_ordinal": ordinal, + "q_projection_ordinal": q_projection.ordinal, + "q_rope_ordinal": int(rope["ordinal"]), + "k_projection_ordinal": k_projection.ordinal, + "k_rope_set_rows_ordinal": int(rope_set_rows["ordinal"]), + "v_projection_ordinal": v_projection.ordinal, + "v_set_rows_ordinal": int(set_rows["ordinal"]), + "postprocess_ns": occurrence_ns, + "head": { + "query_heads": compile_parameters(command).get("ggml.flash_attention.query_head_count"), + "key_value_heads": compile_parameters(command).get("ggml.flash_attention.key_value_head_count"), + "qk_head_size": compile_parameters(command).get("ggml.flash_attention.qk_head_size"), + "value_head_size": compile_parameters(command).get("ggml.flash_attention.value_head_size"), + }, + } + ) + add_candidate( + candidates, + "decode attention projection postprocess", + "attention_qkv_decode", + attention_occurrences, + attention_chain_ns, + attention_removable_ns, + [ + "Current decode path uses standalone Q rope, K rope+cache write, and V cache write before split flash attention.", + "Existing llm_attention_qkv fused projection matchers require token_count > 1, so this looks like decode-route coverage is missing.", + ], + ) + + swiglu_occurrences = [] + swiglu_chain_ns = 0.0 + swiglu_removable_ns = 0.0 + residual_occurrences = [] + residual_chain_ns = 0.0 + residual_removable_ns = 0.0 + tail_occurrences = [] + tail_chain_ns = 0.0 + tail_removable_ns = 0.0 + for command in commands: + if command["kernel"] != "loom_libs:ggml_binary_f32": + continue + ordinal = int(command["ordinal"]) + current_ns = command_metric_ns(command, shape_info) + consumers = output_consumers(command, fanout) + lhs = producer_for_read(command, "lhs", read_producers) + rhs = producer_for_read(command, "rhs", read_producers) + op = op_name(command) + if op == "swiglu": + down_projection = consumers[0] if len(consumers) == 1 else None + occurrence = { + "binary_ordinal": ordinal, + "lhs_producer": lhs.ordinal if lhs else None, + "rhs_producer": rhs.ordinal if rhs else None, + "consumer": down_projection.ordinal if down_projection else None, + "element_count": integer_parameters(command).get("element_count"), + "binary_ns": current_ns, + } + swiglu_occurrences.append(occurrence) + swiglu_removable_ns += current_ns + swiglu_chain_ns += current_ns + if lhs is not None: + swiglu_chain_ns += command_metric_ns(by_ordinal[lhs.ordinal], shape_info) + if rhs is not None: + swiglu_chain_ns += command_metric_ns(by_ordinal[rhs.ordinal], shape_info) + if down_projection is not None: + swiglu_chain_ns += command_metric_ns(by_ordinal[down_projection.ordinal], shape_info) + elif op == "add": + consumer_kernels = [consumer.kernel for consumer in consumers] + occurrence = { + "binary_ordinal": ordinal, + "lhs_producer": lhs.ordinal if lhs else None, + "rhs_producer": rhs.ordinal if rhs else None, + "consumers": [ + { + "ordinal": consumer.ordinal, + "kernel": consumer.kernel, + "binding": consumer.binding, + } + for consumer in consumers + ], + "element_count": integer_parameters(command).get("element_count"), + "binary_ns": current_ns, + } + if "loom_libs:ggml_rmsnorm_binary_f32" in consumer_kernels: + residual_occurrences.append(occurrence) + residual_removable_ns += current_ns + residual_chain_ns += current_ns + for consumer in consumers: + if consumer.kernel == "loom_libs:ggml_rmsnorm_binary_f32": + residual_chain_ns += command_metric_ns(by_ordinal[consumer.ordinal], shape_info) + if any(consumer.kernel == "loom_libs:ggml_rmsnorm_binary_f32" for consumer in consumers) and ( + (lhs and lhs.kernel == "hrx:ggml_gather_add_f32") or (rhs and rhs.kernel == "hrx:ggml_gather_add_f32") + ): + tail_occurrences.append(occurrence) + tail_removable_ns += current_ns + tail_chain_ns += current_ns + for consumer in consumers: + tail_chain_ns += command_metric_ns(by_ordinal[consumer.ordinal], shape_info) + + add_candidate( + candidates, + "decode SwiGLU MLP", + "swiglu_decode", + swiglu_occurrences, + swiglu_chain_ns, + swiglu_removable_ns, + [ + "Pattern is two decode matmuls feeding binary_f32 op=4, then a down projection.", + "This is separate from existing binary_swiglu_symmetric_i4_k32 coverage.", + ], + ) + add_candidate( + candidates, + "residual add next RMS norm", + "add_next_rmsnorm_f32", + residual_occurrences, + residual_chain_ns, + residual_removable_ns, + [ + "Pattern is binary_f32 op=0 feeding rmsnorm_binary_f32.", + "Some add outputs have multiple consumers, so the fused kernel must still publish the residual/add output when it is reused.", + ], + ) + add_candidate( + candidates, + "final output tail add norm", + "output_tail", + tail_occurrences, + tail_chain_ns, + tail_removable_ns, + [ + "Low-count final path involving gather_add, binary add, RMS norm, and output projection.", + ], + ) + + matmul_add_norm_occurrences = [] + matmul_add_norm_chain_ns = 0.0 + for command in commands: + if command["kernel"] != "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32": + continue + lhs = producer_for_read(command, "lhs", read_producers) + if lhs is None or lhs.kernel != "loom_libs:ggml_mul_mat_f32_f32_decode_wave64": + continue + ordinal = int(command["ordinal"]) + current_ns = command_metric_ns(command, shape_info) + lhs_ns = command_metric_ns(by_ordinal[lhs.ordinal], shape_info) + matmul_add_norm_chain_ns += current_ns + lhs_ns + matmul_add_norm_occurrences.append( + { + "add_rmsnorm_ordinal": ordinal, + "matmul_ordinal": lhs.ordinal, + "hidden_size": compile_parameters(command).get("ggml.add_rmsnorm_binary_symmetric_i4.hidden_size"), + "chain_ns": current_ns + lhs_ns, + } + ) + add_candidate( + candidates, + "decode matmul add RMS norm", + "matmul_add_next_rmsnorm_decode", + matmul_add_norm_occurrences, + matmul_add_norm_chain_ns, + 0.0, + [ + "Pattern is decode matmul feeding add_rmsnorm_binary_symmetric_i4.", + "Savings depend on a decode matmul epilogue or fused next-rmsnorm route, so removable standalone time is not estimated here.", + ], + ) + + total_removable = sum(candidate["removable_standalone_ns"] for candidate in candidates) + for candidate in candidates: + candidate["current_chain_ms"] = candidate["current_chain_ns"] / 1_000_000.0 + candidate["removable_standalone_ms"] = candidate["removable_standalone_ns"] / 1_000_000.0 + candidate["removable_share_of_candidates"] = ( + candidate["removable_standalone_ns"] / total_removable * 100.0 if total_removable else 0.0 + ) + return sorted(candidates, key=lambda item: (item["removable_standalone_ns"], item["current_chain_ns"]), reverse=True) + + +def sequential_triples(commands: list[dict[str, Any]], limit: int = 30) -> list[dict[str, Any]]: + counts: Counter[tuple[str, str, str]] = Counter() + ordered = sorted(commands, key=lambda item: int(item["ordinal"])) + for first, second, third in zip(ordered, ordered[1:], ordered[2:], strict=False): + counts[(first["kernel"], second["kernel"], third["kernel"])] += 1 + return [ + {"kernels": list(kernels), "count": count} + for kernels, count in counts.most_common(limit) + ] + + +def build_report(commands: list[dict[str, Any]], + manifest: dict[str, Any], + summary: dict[str, Any]) -> dict[str, Any]: + shape_info, warnings = map_shape_keys_to_dispatches(commands, manifest, summary) + edges, fanout = exact_transient_edges(commands) + candidates = classify_candidates(commands, shape_info, fanout) + total_weighted_ns = float(summary.get("total_weighted_ns", 0.0)) + for candidate in candidates: + candidate["removable_share_of_total_modelled_kernel_time"] = ( + candidate["removable_standalone_ns"] / total_weighted_ns * 100.0 if total_weighted_ns else 0.0 + ) + candidate["current_chain_share_of_total_modelled_kernel_time"] = ( + candidate["current_chain_ns"] / total_weighted_ns * 100.0 if total_weighted_ns else 0.0 + ) + return { + "schema": "ggml-hrx-model-fusion-adjacency-v1", + "model": manifest.get("model"), + "scenario": manifest.get("scenario"), + "command_count": len(commands), + "dispatch_count": manifest.get("dispatch_count"), + "generated_count": manifest.get("generated_count"), + "summary_metric": summary.get("metric_description"), + "total_weighted_ns": total_weighted_ns, + "warnings": warnings, + "candidate_fusions": candidates, + "exact_transient_edge_summary": summarize_edges(edges), + "fanout_summary": summarize_fanout(fanout), + "sequential_triples": sequential_triples(commands), + } + + +def md_escape(text: Any) -> str: + return str(text).replace("|", "\\|") + + +def format_ms(ns: float) -> str: + return f"{ns / 1_000_000.0:.3f}" + + +def write_markdown(path: Path, report: dict[str, Any]) -> None: + lines = [ + "# HRX Model Fusion Adjacency", + "", + f"Model: `{report.get('model')}`", + f"Scenario: `{report.get('scenario')}`", + f"Commands: `{report.get('command_count')}`", + f"Generated dispatch shapes: `{report.get('generated_count')} / {report.get('dispatch_count')}`", + f"Timing metric: `{report.get('summary_metric')}`", + f"Total weighted modeled kernel time: `{format_ms(float(report.get('total_weighted_ns', 0.0)))} ms`", + "", + ] + if report.get("warnings"): + lines.append("## Warnings") + lines.append("") + for warning in report["warnings"]: + lines.append(f"- {warning}") + lines.append("") + + lines.extend( + [ + "## Candidate Fusions", + "", + "| Candidate | Occurrences | Current chain ms | Removable standalone ms | Current chain share | Removable share | Notes |", + "| --- | ---: | ---: | ---: | ---: | ---: | --- |", + ] + ) + for candidate in report["candidate_fusions"]: + notes = "
".join(md_escape(note) for note in candidate.get("notes", [])) + lines.append( + f"| {md_escape(candidate['name'])} | {candidate['occurrence_count']} | " + f"{candidate['current_chain_ms']:.3f} | {candidate['removable_standalone_ms']:.3f} | " + f"{candidate['current_chain_share_of_total_modelled_kernel_time']:.2f}% | " + f"{candidate['removable_share_of_total_modelled_kernel_time']:.2f}% | {notes} |" + ) + + lines.extend( + [ + "", + "## Exact Transient Producer-Consumer Edges", + "", + "| Producer | Consumer | Edges | Example ordinals |", + "| --- | --- | ---: | --- |", + ] + ) + for edge in report["exact_transient_edge_summary"][:30]: + examples = ", ".join( + f"{example['producer_ordinal']}:{example['producer_binding']}->{example['consumer_ordinal']}:{example['consumer_binding']}" + for example in edge.get("examples", [])[:3] + ) + lines.append( + f"| `{edge['producer_kernel']}` | `{edge['consumer_kernel']}` | {edge['edge_count']} | {md_escape(examples)} |" + ) + + lines.extend( + [ + "", + "## Fanout Patterns", + "", + "| Producer | Consumers | Produced values |", + "| --- | --- | ---: |", + ] + ) + for row in report["fanout_summary"][:25]: + consumers = ", ".join(f"`{kernel}`" for kernel in row["consumer_kernels"]) + lines.append(f"| `{row['producer_kernel']}` | {consumers} | {row['producer_count']} |") + + lines.extend( + [ + "", + "## Top Sequential Triples", + "", + "| Kernels | Count |", + "| --- | ---: |", + ] + ) + for triple in report["sequential_triples"][:20]: + kernels = " -> ".join(f"`{kernel}`" for kernel in triple["kernels"]) + lines.append(f"| {kernels} | {triple['count']} |") + + lines.extend( + [ + "", + "## Candidate Examples", + "", + ] + ) + for candidate in report["candidate_fusions"]: + lines.append(f"### {candidate['name']}") + lines.append("") + if not candidate.get("examples"): + lines.append("No matching occurrences.") + lines.append("") + continue + lines.append("```json") + lines.append(json.dumps(candidate["examples"][:3], indent=2, sort_keys=True)) + lines.append("```") + lines.append("") + + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dump-dir", type=Path, required=True, help="HRX command-program dump directory.") + parser.add_argument("--scenario-manifest", type=Path, required=True, help="Generated model scenario manifest.") + parser.add_argument("--summary-json", type=Path, required=True, help="Summarized Loom benchmark JSON.") + parser.add_argument("--output-json", type=Path, required=True, help="Fusion adjacency JSON output.") + parser.add_argument("--output-md", type=Path, required=True, help="Fusion adjacency Markdown output.") + args = parser.parse_args() + + commands = load_commands(args.dump_dir) + manifest = load_json(args.scenario_manifest) + summary = load_json(args.summary_json) + report = build_report(commands, manifest, summary) + write_json(args.output_json, report) + write_markdown(args.output_md, report) + print(f"wrote {args.output_json}") + print(f"wrote {args.output_md}") + + +if __name__ == "__main__": + main() diff --git a/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py new file mode 100755 index 000000000000..4b9084678f91 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/generate-model-benchmarks.py @@ -0,0 +1,1048 @@ +#!/usr/bin/env python3 +# +# Generate model-scoped Loom benchmarks from HRX command program dumps. + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +SCRIPT_DIR = Path(__file__).resolve().parent +HRX_DIR = SCRIPT_DIR.parents[1] +BENCHMARK_DIR = HRX_DIR / "benchmarks" +KERNEL_CORPUS_DIR = HRX_DIR / "kernel-corpus" / "kernels" +KERNEL_TARGET_RE = re.compile( + r"kernel\.def\s+target\(@(?P[A-Za-z0-9_.$-]+)\)" + r"(?:\s+export\(\"[^\"]+\"\))?\s+@(?P[A-Za-z0-9_.$-]+)\s*\(" +) + + +@dataclass +class ScenarioBenchmarks: + scenario: str + command_count: int + dispatch_count: int + generated_count: int + kernel_counts: dict[str, int] + dispatches: list[dict[str, Any]] + cases: list[str] + used_exports: dict[str, dict[str, Any]] + + +def fail(message: str) -> None: + raise SystemExit(message) + + +def load_json(path: Path) -> Any: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def read_text(path: Path) -> str: + return path.read_text(encoding="utf-8") + + +def write_json(path: Path, data: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, indent=2, sort_keys=True) + f.write("\n") + + +def sanitize_symbol(text: str) -> str: + text = text.split(":")[-1] + text = re.sub(r"[^A-Za-z0-9_]+", "_", text) + text = re.sub(r"_+", "_", text).strip("_") + if not text: + return "unknown" + if text[0].isdigit(): + return "k_" + text + return text + + +def int_param(command: dict[str, Any], name: str, default: int | None = None) -> int: + params = command.get("integer_parameters", {}) + if name in params: + return int(params[name]) + if default is not None: + return default + fail(f"command {command.get('ordinal')} {command.get('kernel')} is missing integer parameter {name}") + + +def config_int(command: dict[str, Any], name: str, default: int | None = None) -> int: + params = command.get("compile_parameters", {}) + if name in params: + return int(params[name]) + if default is not None: + return default + fail(f"command {command.get('ordinal')} {command.get('kernel')} is missing compile parameter {name}") + + +def binding_length(command: dict[str, Any], name: str) -> int: + for binding in command.get("bindings", []): + if binding.get("name") == name: + return int(binding["length"]) + fail(f"command {command.get('ordinal')} {command.get('kernel')} is missing binding {name}") + + +def ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def tensor_type_for_format(format_value: int) -> str | None: + if format_value == 16: + return "f16" + if format_value == 30: + return "bf16" + if format_value == 32: + return "f32" + return None + + +def fill_tensor(name: str, value: str, shape: str, dtype: str) -> str: + return f" %{name} = check.generate.fill value({value}) : tensor<{shape}x{dtype}>" + + +def iota_tensor(name: str, offset: str, step: str, shape: str, dtype: str, period: int | None = None) -> str: + period_text = "" if period is None else f" period({period})" + return f" %{name} = check.generate.iota offset({offset}) step({step}){period_text} : tensor<{shape}x{dtype}>" + + +def case_binary(symbol: str, command: dict[str, Any]) -> str: + count = int_param(command, "element_count") + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %element_count = check.literal value({count}) : index", + fill_tensor("lhs", "2.0", f"{count}", "f32"), + fill_tensor("rhs", "3.0", f"{count}", "f32"), + fill_tensor("output", "0.0", f"{count}", "f32"), + f" kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<{count}xf32>, tensor<{count}xf32>, tensor<{count}xf32>)", + " check.return", + "}", + ] + ) + + +def case_rmsnorm_binary(symbol: str, command: dict[str, Any]) -> str: + token_count = int_param(command, "token_count") + hidden_size = config_int(command, "ggml.rmsnorm_binary_f32.hidden_size") + shape = f"{token_count}x{hidden_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + fill_tensor("input", "2.0", shape, "f32"), + fill_tensor("rhs", "3.0", f"{hidden_size}", "f32"), + fill_tensor("output", "0.0", shape, "f32"), + f" kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<{shape}xf32>, tensor<{hidden_size}xf32>, tensor<{shape}xf32>)", + " check.return", + "}", + ] + ) + + +def case_mul_mat(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + if kernel == "ggml_mul_mat_f32_f32_decode_wave64": + input_size = int_param(command, "input_size") + output_size = int_param(command, "output_size") + accumulation = 0 + weight_format = config_int(command, "ggml.mul_mat_f32_f32_decode.weight_format") + else: + input_size = config_int(command, "ggml.mul_mat.input_size") + output_size = config_int(command, "ggml.mul_mat.output_size") + accumulation = config_int(command, "ggml.mul_mat.output_accumulation", 0) + weight_format = config_int(command, "ggml.mul_mat.weight_format") + weight_type = tensor_type_for_format(weight_format) + if weight_type is None: + return None + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + shape_out = f"{token_count}x{output_size}" + output_init = 1.0 if accumulation else 0.0 + if kernel == "ggml_mul_mat_f32_f32_decode_wave64": + launch_params = "%token_count, %input_size, %output_size" + workload_params = "%token_count, %input_size, %output_size" + workload_types = "[index, index, index]" + launch_types = "index, index, index" + literals = [ + f" %token_count = check.literal value({token_count}) : index", + f" %input_size = check.literal value({input_size}) : index", + f" %output_size = check.literal value({output_size}) : index", + ] + else: + launch_params = "%token_count" + workload_params = "%token_count" + workload_types = "[index]" + launch_types = "index" + literals = [f" %token_count = check.literal value({token_count}) : index"] + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + *literals, + fill_tensor("input", "1.0", shape_in, "f32"), + fill_tensor("weight", "1.0", shape_weight, weight_type), + fill_tensor("output", f"{output_init:.1f}", shape_out, "f32"), + f" kernel.launch @{kernel}[{workload_params}]({launch_params}, %input, %weight, %output) : {workload_types}({launch_types}, tensor<{shape_in}xf32>, tensor<{shape_weight}x{weight_type}>, tensor<{shape_out}xf32>)", + " check.return", + "}", + ] + ) + + +def case_mul_mat_add_decode(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + input_size = int_param(command, "input_size") + output_size = int_param(command, "output_size") + weight_format = config_int(command, "ggml.mul_mat_f32_f32_decode.weight_format") + weight_type = tensor_type_for_format(weight_format) + if weight_type is None: + return None + + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + shape_out = f"{token_count}x{output_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + f" %input_size = check.literal value({input_size}) : index", + f" %output_size = check.literal value({output_size}) : index", + fill_tensor("input", "1.0", shape_in, "f32"), + fill_tensor("weight", "1.0", shape_weight, weight_type), + fill_tensor("residual_input", "0.25", shape_out, "f32"), + fill_tensor("residual_output", "0.0", shape_out, "f32"), + f" kernel.launch @{kernel}[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %residual_input, %residual_output) : [index, index, index](index, index, index, tensor<{shape_in}xf32>, tensor<{shape_weight}x{weight_type}>, tensor<{shape_out}xf32>, tensor<{shape_out}xf32>)", + " check.return", + "}", + ] + ) + + +def case_mul_mat_postops(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + input_size = config_int(command, "ggml.mul_mat_postops.input_size") + output_size = config_int(command, "ggml.mul_mat_postops.output_size") + weight_format = config_int(command, "ggml.mul_mat_postops.weight_format") + weight_type = tensor_type_for_format(weight_format) + if weight_type is None: + return None + + has_bias = "bias" in kernel + has_residual = "add" in kernel + has_rmsnorm = "next_rmsnorm" in kernel + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + shape_out = f"{token_count}x{output_size}" + completion_count = (token_count + 31) // 32 + + tensors = [ + fill_tensor("input", "1.0", shape_in, "f32"), + fill_tensor("weight", "1.0", shape_weight, weight_type), + ] + launch_values = ["%token_count", "%input", "%weight"] + launch_types = ["index", f"tensor<{shape_in}xf32>", f"tensor<{shape_weight}x{weight_type}>"] + + if has_bias: + tensors.append(fill_tensor("bias", "0.5", f"{output_size}", "f32")) + launch_values.append("%bias") + launch_types.append(f"tensor<{output_size}xf32>") + if has_residual: + tensors.append(fill_tensor("residual_input", "0.25", shape_out, "f32")) + tensors.append(fill_tensor("residual_output", "0.0", shape_out, "f32")) + launch_values.extend(["%residual_input", "%residual_output"]) + launch_types.extend([f"tensor<{shape_out}xf32>", f"tensor<{shape_out}xf32>"]) + else: + tensors.append(fill_tensor("output", "0.0", shape_out, "f32")) + launch_values.append("%output") + launch_types.append(f"tensor<{shape_out}xf32>") + + if has_rmsnorm: + tensors.append(fill_tensor("norm_weight", "1.0", f"{output_size}", "f32")) + tensors.append(fill_tensor("normalized_output", "0.0", shape_out, "f32")) + tensors.append(fill_tensor("completion_counters", "0", f"{completion_count}", "i32")) + launch_values.extend(["%norm_weight", "%normalized_output", "%completion_counters"]) + launch_types.extend([f"tensor<{output_size}xf32>", f"tensor<{shape_out}xf32>", f"tensor<{completion_count}xi32>"]) + + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + *tensors, + f" kernel.launch @{kernel}[%token_count]({', '.join(launch_values)}) : [index]({', '.join(launch_types)})", + " check.return", + "}", + ] + ) + + +def case_swiglu(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + input_size = config_int(command, "ggml.mul_mat_swiglu.input_size") + output_size = config_int(command, "ggml.mul_mat_swiglu.output_size") + gate_format = config_int(command, "ggml.mul_mat_swiglu.gate_weight_format") + up_format = config_int(command, "ggml.mul_mat_swiglu.up_weight_format") + gate_type = tensor_type_for_format(gate_format) + up_type = tensor_type_for_format(up_format) + if gate_type is None or up_type is None: + return None + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + shape_out = f"{token_count}x{output_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + fill_tensor("input", "0.0", shape_in, "f32"), + fill_tensor("gate_weight", "1.0", shape_weight, gate_type), + fill_tensor("up_weight", "1.0", shape_weight, up_type), + fill_tensor("output", "1.0", shape_out, "f32"), + f" kernel.launch @{kernel}[%token_count](%token_count, %input, %gate_weight, %up_weight, %output) : [index](index, tensor<{shape_in}xf32>, tensor<{shape_weight}x{gate_type}>, tensor<{shape_weight}x{up_type}>, tensor<{shape_out}xf32>)", + " check.return", + "}", + ] + ) + + +def case_llm_attention_q_matmul_rope(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + input_size = config_int(command, "llm.attention_qkv.input_size") + output_size = config_int(command, "llm.attention_qkv.output_size") + weight_format = config_int(command, "llm.attention_qkv.weight_format") + head_size = config_int(command, "llm.attention_qkv.head_size") + head_count = config_int(command, "llm.attention_qkv.head_count") + weight_type = tensor_type_for_format(weight_format) + if weight_type is None: + return None + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + shape_out = f"{token_count}x{head_count}x{head_size}" + half_head_size = head_size // 2 + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + fill_tensor("input", "1.0", shape_in, "f32"), + fill_tensor("weight", "1.0", shape_weight, weight_type), + fill_tensor("positions", "0", f"{token_count}", "i32"), + fill_tensor("theta", "0.0", f"{half_head_size}", "f32"), + fill_tensor("freq_factors", "0.0", f"{half_head_size}", "f32"), + fill_tensor("output", "0.0", shape_out, "f32"), + f" kernel.launch @{kernel}[%token_count](%token_count, %input, %weight, %positions, %theta, %freq_factors, %output) : [index](index, tensor<{shape_in}xf32>, tensor<{shape_weight}x{weight_type}>, tensor<{token_count}xi32>, tensor<{half_head_size}xf32>, tensor<{half_head_size}xf32>, tensor<{shape_out}xf32>)", + " check.return", + "}", + ] + ) + + +def case_llm_attention_k_matmul_rope_set_rows(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + input_size = config_int(command, "llm.attention_qkv.input_size") + output_size = config_int(command, "llm.attention_qkv.output_size") + weight_format = config_int(command, "llm.attention_qkv.weight_format") + head_size = config_int(command, "llm.attention_qkv.head_size") + head_count = config_int(command, "llm.attention_qkv.head_count") + cache_rows = config_int(command, "llm.attention_qkv.cache_row_count") + cache_format = config_int(command, "llm.attention_qkv.cache_output_format") + weight_type = tensor_type_for_format(weight_format) + cache_type = tensor_type_for_format(cache_format) + if weight_type is None or cache_type is None: + return None + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + cache_shape = f"{cache_rows}x{head_count}x{head_size}" + half_head_size = head_size // 2 + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + fill_tensor("input", "1.0", shape_in, "f32"), + fill_tensor("weight", "1.0", shape_weight, weight_type), + fill_tensor("positions", "0", f"{token_count}", "i32"), + iota_tensor("indices", "0", "1", f"{token_count}", "i64", period=cache_rows), + fill_tensor("theta", "0.0", f"{half_head_size}", "f32"), + fill_tensor("freq_factors", "0.0", f"{half_head_size}", "f32"), + fill_tensor("cache", "0.0", cache_shape, cache_type), + f" kernel.launch @{kernel}[%token_count](%token_count, %input, %weight, %positions, %indices, %theta, %freq_factors, %cache) : [index](index, tensor<{shape_in}xf32>, tensor<{shape_weight}x{weight_type}>, tensor<{token_count}xi32>, tensor<{token_count}xi64>, tensor<{half_head_size}xf32>, tensor<{half_head_size}xf32>, tensor<{cache_shape}x{cache_type}>)", + " check.return", + "}", + ] + ) + + +def case_llm_attention_v_matmul_set_rows(symbol: str, command: dict[str, Any], kernel: str) -> str | None: + token_count = int_param(command, "token_count") + input_size = config_int(command, "llm.attention_qkv.input_size") + output_size = config_int(command, "llm.attention_qkv.output_size") + weight_format = config_int(command, "llm.attention_qkv.weight_format") + cache_rows = config_int(command, "llm.attention_qkv.cache_row_count") + cache_format = config_int(command, "llm.attention_qkv.cache_output_format") + weight_type = tensor_type_for_format(weight_format) + cache_type = tensor_type_for_format(cache_format) + if weight_type is None or cache_type is None: + return None + shape_in = f"{token_count}x{input_size}" + shape_weight = f"{output_size}x{input_size}" + cache_shape = f"{cache_rows}x{output_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + fill_tensor("input", "1.0", shape_in, "f32"), + fill_tensor("weight", "1.0", shape_weight, weight_type), + iota_tensor("indices", "0", "1", f"{token_count}", "i64", period=cache_rows), + fill_tensor("cache", "0.0", cache_shape, cache_type), + f" kernel.launch @{kernel}[%token_count](%token_count, %input, %weight, %indices, %cache) : [index](index, tensor<{shape_in}xf32>, tensor<{shape_weight}x{weight_type}>, tensor<{token_count}xi64>, tensor<{cache_shape}x{cache_type}>)", + " check.return", + "}", + ] + ) + + +def case_flash_attention(symbol: str, command: dict[str, Any], kernel: str, prefix: str) -> str: + query_tokens = int_param(command, "query_token_count") + kv_tokens = int_param(command, "key_value_token_count") + if prefix == "ggml": + query_heads = config_int(command, "ggml.flash_attention.query_head_count") + kv_heads = config_int(command, "ggml.flash_attention.key_value_head_count") + legacy_head_size = config_int(command, "ggml.flash_attention.head_size", 0) + qk_head_size = config_int(command, "ggml.flash_attention.qk_head_size", legacy_head_size) + value_head_size = config_int(command, "ggml.flash_attention.value_head_size", qk_head_size) + else: + query_heads = config_int(command, f"{prefix}.attention.query_head_count") + kv_heads = config_int(command, f"{prefix}.attention.key_value_head_count") + qk_head_size = 128 + value_head_size = 128 + query_shape = f"{query_tokens}x{query_heads}x{qk_head_size}" + key_shape = f"{kv_tokens}x{kv_heads}x{qk_head_size}" + value_shape = f"{kv_tokens}x{kv_heads}x{value_head_size}" + output_shape = f"{query_tokens}x{query_heads}x{value_head_size}" + mask_shape = f"{query_tokens}x{kv_tokens}" + if prefix == "ggml": + gate_bytes = binding_length(command, "gate") + if gate_bytes % 4 != 0: + fail( + f"command {command.get('ordinal')} {command.get('kernel')} has " + f"non-f32 gate binding length {gate_bytes}" + ) + gate_elements = gate_bytes // 4 + apply_gate = config_int(command, "ggml.flash_attention.apply_gate", 0) + if apply_gate: + gate_stride_head = config_int(command, "ggml.flash_attention.gate_stride_head") + gate_stride_token = config_int(command, "ggml.flash_attention.gate_stride_token") + required_gate_elements = ( + (query_tokens - 1) * gate_stride_token + + (query_heads - 1) * gate_stride_head + + value_head_size + ) + if gate_elements < required_gate_elements: + fail( + f"command {command.get('ordinal')} {command.get('kernel')} gate binding " + f"has {gate_elements} f32 elements but the strided access requires " + f"{required_gate_elements}" + ) + launch = ( + f" kernel.launch @{kernel}[%query_token_count, %key_value_token_count]" + f"(%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %gate, %output)" + f" : [index, index](index, index, tensor<{query_shape}xf32>, tensor<{key_shape}xf16>, " + f"tensor<{value_shape}xf16>, tensor<{mask_shape}xf16>, tensor<{gate_elements}xf32>, " + f"tensor<{output_shape}xf32>)" + ) + else: + launch = ( + f" kernel.launch @{kernel}[%query_token_count, %key_value_token_count]" + f"(%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output)" + f" : [index, index](index, index, tensor<{query_shape}xf32>, tensor<{key_shape}xf16>, " + f"tensor<{value_shape}xf16>, tensor<{mask_shape}xf16>, tensor<{output_shape}xf32>)" + ) + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %query_token_count = check.literal value({query_tokens}) : index", + f" %key_value_token_count = check.literal value({kv_tokens}) : index", + fill_tensor("query", "0.0", query_shape, "f32"), + fill_tensor("key", "0.0", key_shape, "f16"), + fill_tensor("value", "0.0", value_shape, "f16"), + fill_tensor("mask", "0.0", mask_shape, "f16"), + *([fill_tensor("gate", "0.0", f"{gate_elements}", "f32")] if prefix == "ggml" else []), + fill_tensor("output", "1.0", output_shape, "f32"), + launch, + " check.return", + "}", + ] + ) + + +def case_flash_attention_decode_split(symbol: str, command: dict[str, Any], kernel: str) -> str: + kv_tokens = int_param(command, "key_value_token_count") + kv_capacity = config_int(command, "ggml.flash_attention.decode.key_value_token_capacity") + query_heads = config_int(command, "ggml.flash_attention.query_head_count") + kv_heads = config_int(command, "ggml.flash_attention.key_value_head_count") + legacy_head_size = config_int(command, "ggml.flash_attention.head_size", 0) + qk_head_size = config_int(command, "ggml.flash_attention.qk_head_size", legacy_head_size) + value_head_size = config_int(command, "ggml.flash_attention.value_head_size", qk_head_size) + partial_blocks = ceil_div(kv_capacity, 64) + query_shape = f"{query_heads}x{qk_head_size}" + key_shape = f"{kv_tokens}x{kv_heads}x{qk_head_size}" + value_shape = f"{kv_tokens}x{kv_heads}x{value_head_size}" + mask_shape = f"{kv_tokens}" + partial_shape = f"{kv_heads}x{partial_blocks}x16" + partial_output_shape = f"{kv_heads}x{partial_blocks}x16x{value_head_size}" + completion_shape = f"{kv_heads}" + output_shape = f"{query_heads}x{value_head_size}" + tensors = [ + fill_tensor("query", "0.0", query_shape, "f32"), + fill_tensor("key", "0.0", key_shape, "f16"), + fill_tensor("value", "0.0", value_shape, "f16"), + fill_tensor("mask", "0.0", mask_shape, "f16"), + fill_tensor("partial_max", "0.0", partial_shape, "f32"), + fill_tensor("partial_sum", "0.0", partial_shape, "f32"), + fill_tensor("partial_output", "0.0", partial_output_shape, "f16"), + fill_tensor("completion_counter", "0", completion_shape, "i32"), + fill_tensor("output", "1.0", output_shape, "f32"), + ] + launch_values = [ + "%key_value_token_count", + "%query", + "%key", + "%value", + "%mask", + "%partial_max", + "%partial_sum", + "%partial_output", + "%completion_counter", + "%output", + ] + launch_types = [ + "index", + f"tensor<{query_shape}xf32>", + f"tensor<{key_shape}xf16>", + f"tensor<{value_shape}xf16>", + f"tensor<{mask_shape}xf16>", + f"tensor<{partial_shape}xf32>", + f"tensor<{partial_shape}xf32>", + f"tensor<{partial_output_shape}xf16>", + f"tensor<{completion_shape}xi32>", + f"tensor<{output_shape}xf32>", + ] + if kernel.endswith("_next_q8"): + next_q8_bytes = binding_length(command, "next_q8_output") + tensors.append(fill_tensor("next_q8_output", "0", f"{next_q8_bytes}", "i8")) + launch_values.append("%next_q8_output") + launch_types.append(f"tensor<{next_q8_bytes}xi8>") + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %key_value_token_count = check.literal value({kv_tokens}) : index", + *tensors, + f" kernel.launch @{kernel}[%key_value_token_count]({', '.join(launch_values)}) : [index]({', '.join(launch_types)})", + " check.return", + "}", + ] + ) + + +def case_rope(symbol: str, command: dict[str, Any], kernel: str, prefix: str) -> str: + token_count = int_param(command, "token_count") + head_size = config_int(command, f"{prefix}.head_size") + head_count = config_int(command, f"{prefix}.head_count") + half_head_size = head_size // 2 + data_shape = f"{token_count}x{head_count}x{head_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + fill_tensor("positions", "0", f"{token_count}", "i32"), + fill_tensor("input", "0.0", data_shape, "f32"), + fill_tensor("theta", "0.0", f"{half_head_size}", "f32"), + fill_tensor("freq_factors", "0.0", f"{half_head_size}", "f32"), + fill_tensor("output", "1.0", data_shape, "f32"), + f" kernel.launch @{kernel}[%token_count](%token_count, %positions, %input, %theta, %freq_factors, %output) : [index](index, tensor<{token_count}xi32>, tensor<{data_shape}xf32>, tensor<{half_head_size}xf32>, tensor<{half_head_size}xf32>, tensor<{data_shape}xf32>)", + " check.return", + "}", + ] + ) + + +def case_rope_set_rows(symbol: str, command: dict[str, Any]) -> str | None: + token_count = int_param(command, "token_count") + cache_rows = int_param(command, "cache_row_count") + head_size = config_int(command, "ggml.rope_set_rows_f32.head_size") + head_count = config_int(command, "ggml.rope_set_rows_f32.head_count") + output_format = config_int(command, "ggml.rope_set_rows_f32.output_format") + output_type = tensor_type_for_format(output_format) + if output_type is None: + return None + half_head_size = head_size // 2 + input_shape = f"{token_count}x{head_count}x{head_size}" + cache_shape = f"{cache_rows}x{head_count}x{head_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + f" %cache_row_count = check.literal value({cache_rows}) : index", + fill_tensor("positions", "0", f"{token_count}", "i32"), + iota_tensor("indices", "0", "1", f"{token_count}", "i64", period=cache_rows), + fill_tensor("input", "0.0", input_shape, "f32"), + fill_tensor("theta", "0.0", f"{half_head_size}", "f32"), + fill_tensor("freq_factors", "0.0", f"{half_head_size}", "f32"), + fill_tensor("cache", "0.0", cache_shape, output_type), + f" kernel.launch @ggml_rope_set_rows_f32[%token_count, %cache_row_count](%token_count, %cache_row_count, %positions, %indices, %input, %theta, %freq_factors, %cache) : [index, index](index, index, tensor<{token_count}xi32>, tensor<{token_count}xi64>, tensor<{input_shape}xf32>, tensor<{half_head_size}xf32>, tensor<{half_head_size}xf32>, tensor<{cache_shape}x{output_type}>)", + " check.return", + "}", + ] + ) + + +def case_set_rows(symbol: str, command: dict[str, Any]) -> str | None: + token_count = int_param(command, "token_count") + cache_rows = int_param(command, "cache_row_count") + hidden_size = int_param(command, "hidden_size") + input_format = config_int(command, "ggml.set_rows.input_format") + output_format = config_int(command, "ggml.set_rows.output_format") + input_type = tensor_type_for_format(input_format) + output_type = tensor_type_for_format(output_format) + if input_type is None or output_type is None: + return None + rows_shape = f"{token_count}x{hidden_size}" + cache_shape = f"{cache_rows}x{hidden_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + f" %cache_row_count = check.literal value({cache_rows}) : index", + f" %hidden_size = check.literal value({hidden_size}) : index", + fill_tensor("rows", "0.0", rows_shape, input_type), + iota_tensor("indices", "0", "1", f"{token_count}", "i64", period=cache_rows), + fill_tensor("cache", "0.0", cache_shape, output_type), + f" kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<{rows_shape}x{input_type}>, tensor<{token_count}xi64>, tensor<{cache_shape}x{output_type}>)", + " check.return", + "}", + ] + ) + + +def case_get_rows(symbol: str, command: dict[str, Any]) -> str | None: + token_count = int_param(command, "token_count") + row_count = int_param(command, "row_count") + hidden_size = int_param(command, "hidden_size") + weight_format = config_int(command, "ggml.get_rows_f32.weight_format") + weight_type = tensor_type_for_format(weight_format) + if weight_type is None: + return None + output_shape = f"{token_count}x{hidden_size}" + weight_shape = f"{row_count}x{hidden_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %token_count = check.literal value({token_count}) : index", + f" %row_count = check.literal value({row_count}) : index", + f" %hidden_size = check.literal value({hidden_size}) : index", + fill_tensor("token_ids", "0", f"{token_count}", "i32"), + fill_tensor("weight", "0.0", weight_shape, weight_type), + fill_tensor("output", "1.0", output_shape, "f32"), + f" kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<{token_count}xi32>, tensor<{weight_shape}x{weight_type}>, tensor<{output_shape}xf32>)", + " check.return", + "}", + ] + ) + + +def case_gather_add(symbol: str, command: dict[str, Any]) -> str: + source_tokens = int_param(command, "source_token_count") + output_tokens = int_param(command, "output_token_count") + hidden_size = int_param(command, "hidden_size") + source_shape = f"{source_tokens}x{hidden_size}" + output_shape = f"{output_tokens}x{hidden_size}" + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + f" %source_token_count = check.literal value({source_tokens}) : index", + f" %output_token_count = check.literal value({output_tokens}) : index", + f" %hidden_size = check.literal value({hidden_size}) : index", + fill_tensor("attention", "0.0", source_shape, "f32"), + fill_tensor("residual", "0.0", source_shape, "f32"), + fill_tensor("output_ids", "0", f"{output_tokens}", "i32"), + fill_tensor("output", "1.0", output_shape, "f32"), + f" kernel.launch @ggml_gather_add_f32[%source_token_count, %output_token_count, %hidden_size](%source_token_count, %output_token_count, %hidden_size, %attention, %residual, %output_ids, %output) : [index, index, index](index, index, index, tensor<{source_shape}xf32>, tensor<{source_shape}xf32>, tensor<{output_tokens}xi32>, tensor<{output_shape}xf32>)", + " check.return", + "}", + ] + ) + + +def case_generic(symbol: str, command: dict[str, Any], export: dict[str, Any]) -> str | None: + integer_parameters = command.get("integer_parameters", {}) + binding_lengths = { + binding.get("name"): int(binding["length"]) + for binding in command.get("bindings", []) + if binding.get("name") is not None and "length" in binding + } + workload_parameters = list(export.get("workload_parameters", [])) + launch_parameters = list(export.get("launch_parameters", [])) + bindings = list(export.get("bindings", [])) + + scalar_names: list[str] = [] + for parameter in [*workload_parameters, *launch_parameters]: + name = parameter["name"] + if name not in scalar_names: + scalar_names.append(name) + missing_scalars = [name for name in scalar_names if name not in integer_parameters] + missing_bindings = [name for name in bindings if name not in binding_lengths] + if missing_scalars or missing_bindings: + return None + + literals = [ + f" %{name} = check.literal value({int(integer_parameters[name])}) : index" + for name in scalar_names + ] + tensors = [ + fill_tensor(name, "0", f"{binding_lengths[name]}", "i8") + for name in bindings + ] + workload_values = ", ".join(f"%{parameter['name']}" for parameter in workload_parameters) + launch_values = [f"%{parameter['name']}" for parameter in launch_parameters] + launch_values.extend(f"%{name}" for name in bindings) + workload_types = ", ".join("index" for _ in workload_parameters) + launch_types = ["index" for _ in launch_parameters] + launch_types.extend(f"tensor<{binding_lengths[name]}xi8>" for name in bindings) + return "\n".join( + [ + f"check.case public @{symbol}_case {{", + *literals, + *tensors, + f" kernel.launch @{export['symbol']}[{workload_values}]({', '.join(launch_values)}) : [{workload_types}]({', '.join(launch_types)})", + " check.return", + "}", + ] + ) + + +def render_case(symbol: str, command: dict[str, Any], export: dict[str, Any]) -> str | None: + kernel = command["kernel"].split(":")[-1] + case: str | None + if kernel == "ggml_binary_f32": + case = case_binary(symbol, command) + elif kernel == "ggml_rmsnorm_binary_f32": + case = case_rmsnorm_binary(symbol, command) + elif kernel in ("ggml_mul_mat_f32_f32_wmma", "ggml_mul_mat_f32_f32_decode_wave64"): + case = case_mul_mat(symbol, command, kernel) + elif kernel == "ggml_mul_mat_add_f32_f32_decode_wave64": + case = case_mul_mat_add_decode(symbol, command, kernel) + elif kernel in ( + "ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + "ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + ): + case = case_mul_mat_postops(symbol, command, kernel) + elif kernel in ("ggml_mul_mat_swiglu_f32_f32_wmma", "ggml_mul_mat_swiglu_f32_f32_decode_wave64"): + case = case_swiglu(symbol, command, kernel) + elif kernel in ("llm_attention_q_matmul_rope_f32_f32_wmma", "llm_attention_q_matmul_rope_decode_f32_f32"): + case = case_llm_attention_q_matmul_rope(symbol, command, kernel) + elif kernel in ( + "llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", + "llm_attention_k_matmul_rope_set_rows_decode_f32_f32", + ): + case = case_llm_attention_k_matmul_rope_set_rows(symbol, command, kernel) + elif kernel in ("llm_attention_v_matmul_set_rows_f32_f32_wmma", "llm_attention_v_matmul_set_rows_decode_f32_f32"): + case = case_llm_attention_v_matmul_set_rows(symbol, command, kernel) + elif kernel == "ggml_flash_attention_f32_f16_wmma": + case = case_flash_attention(symbol, command, kernel, "ggml") + elif kernel in ("ggml_flash_attention_decode_split_f32_f16_wmma", "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"): + case = case_flash_attention_decode_split(symbol, command, kernel) + elif kernel == "ggml_rope_f32": + case = case_rope(symbol, command, kernel, "ggml.rope_f32") + elif kernel == "ggml_rope_set_rows_f32": + case = case_rope_set_rows(symbol, command) + elif kernel == "ggml_set_rows": + case = case_set_rows(symbol, command) + elif kernel == "ggml_get_rows_f32": + case = case_get_rows(symbol, command) + elif kernel == "ggml_gather_add_f32": + case = case_gather_add(symbol, command) + else: + case = None + if case is not None: + return case + return case_generic(symbol, command, export) + + +def loom_target_for_export(corpus_dir: Path, export: dict[str, Any]) -> str: + source = corpus_dir / export["source"] + if not source.is_file(): + return "" + text = read_text(source) + symbol = export["symbol"] + for match in KERNEL_TARGET_RE.finditer(text): + if match.group("symbol") == symbol: + return match.group("target") + return "" + + +def load_exports() -> dict[str, dict[str, Any]]: + exports: dict[str, dict[str, Any]] = {} + for manifest_path in sorted(KERNEL_CORPUS_DIR.glob("*/manifest.json")): + manifest = load_json(manifest_path) + corpus_dir = manifest_path.parent + for export in manifest.get("exports", []): + name = export.get("name") + if not name: + continue + export = dict(export) + export["corpus_dir"] = str(corpus_dir.relative_to(HRX_DIR)) + export["loom_target"] = loom_target_for_export(corpus_dir, export) + previous = exports.get(name) + if previous is None or previous.get("target_selector") and not export.get("target_selector"): + exports[name] = export + return exports + + +def render_decl(export: dict[str, Any]) -> str: + symbol = export["symbol"] + workload = ", ".join(f"%{param['name']}: {param['type']}" for param in export.get("workload_parameters", [])) + launch_items = [f"%{param['name']}: {param['type']}" for param in export.get("launch_parameters", [])] + launch_items += [f"%{name}: buffer" for name in export.get("bindings", [])] + launch = ", ".join(launch_items) + target = export.get("loom_target", "") + target_attr = "" if not target else f" target(@{target})" + return f"kernel.decl{target_attr} @{symbol}({workload}) launch({launch})" + + +def command_shape_key(command: dict[str, Any]) -> str: + shape_data = { + "kernel": command.get("kernel"), + "integer_parameters": command.get("integer_parameters", {}), + "compile_parameters": command.get("compile_parameters", {}), + "binding_lengths": [binding.get("length") for binding in command.get("bindings", [])], + } + encoded = json.dumps(shape_data, sort_keys=True, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest()[:16] + + +def load_program_commands(dump_dir: Path) -> list[dict[str, Any]]: + program_paths = sorted(dump_dir.glob("program-*/program.json")) + if not program_paths: + fail(f"no program.json files found under {dump_dir}") + commands: list[dict[str, Any]] = [] + for program_path in program_paths: + program = load_json(program_path) + for command in program.get("commands", []): + command = dict(command) + command["program"] = { + "directory": program_path.parent.name, + "dump_id": program.get("dump_id"), + "shape_hash": program.get("shape_hash"), + "target": program.get("target"), + } + commands.append(command) + return commands + + +def source_paths_for_export(export: dict[str, Any]) -> list[str]: + recipe = export.get("compile_recipe", {}) + paths = [] + for source in recipe.get("primary_sources", []): + paths.append(source) + for source in recipe.get("library_sources", []): + paths.append(source) + return list(dict.fromkeys(paths)) + + +def primary_source_paths_for_export(export: dict[str, Any]) -> list[str]: + return list(dict.fromkeys(export.get("compile_recipe", {}).get("primary_sources", []))) + + +def library_source_paths_for_export(export: dict[str, Any]) -> list[str]: + return list(dict.fromkeys(export.get("compile_recipe", {}).get("library_sources", []))) + + +def generate_scenario(model_slug: str, + scenario: str, + commands: list[dict[str, Any]], + exports: dict[str, dict[str, Any]]) -> ScenarioBenchmarks: + dispatch_entries: list[dict[str, Any]] = [] + cases: list[str] = [] + used_exports: dict[str, dict[str, Any]] = {} + shape_keys = [command_shape_key(command) for command in commands] + shape_counts = Counter(shape_keys) + shape_commands: dict[str, list[dict[str, Any]]] = {} + for command, shape_key in zip(commands, shape_keys, strict=True): + shape_commands.setdefault(shape_key, []).append(command) + per_kernel_counts = Counter(command["kernel"] for command in commands) + + for shape_key, represented_commands in shape_commands.items(): + command = represented_commands[0] + kernel_name = command["kernel"] + export_name = kernel_name.split(":")[-1] + export = exports.get(export_name) + if export is None: + continue + else: + used_exports[export_name] = export + symbol = f"{model_slug}_{scenario}_{len(dispatch_entries):03d}_{sanitize_symbol(export_name)}" + try: + case = render_case(symbol, command, export) + except SystemExit as exc: + case = None + print(f"skipped {kernel_name}: {exc}") + else: + if case is None: + print(f"skipped {kernel_name}: unsupported template") + if case is not None: + cases.append(case) + benchmark_symbol = "@" + symbol + else: + continue + + dispatch_entries.append( + { + "benchmark": benchmark_symbol, + "kernel": kernel_name, + "symbol": export["symbol"], + "count": shape_counts[shape_key], + "integer_parameters": command.get("integer_parameters", {}), + "compile_parameters": command.get("compile_parameters", {}), + "workload_parameters": export.get("workload_parameters", []), + "sources": source_paths_for_export(export), + "primary_sources": primary_source_paths_for_export(export), + "library_sources": library_source_paths_for_export(export), + "corpus_dir": export.get("corpus_dir"), + } + ) + + return ScenarioBenchmarks( + scenario=scenario, + command_count=len(commands), + dispatch_count=len(shape_commands), + generated_count=len(dispatch_entries), + kernel_counts=dict(sorted(per_kernel_counts.items())), + dispatches=dispatch_entries, + cases=cases, + used_exports=used_exports, + ) + + +def write_model_loom(model_file: Path, + model_slug: str, + scenarios: list[ScenarioBenchmarks], + used_exports: dict[str, dict[str, Any]]) -> None: + scenario_names = ", ".join(scenario.scenario for scenario in scenarios) + targets = sorted({ + export["loom_target"] + for export in used_exports.values() + if export.get("loom_target") + }) + decls = [render_decl(export) for _, export in sorted(used_exports.items())] + lines = [ + f"// Generated by tools/benchmarks/generate-model-benchmarks.py for {model_slug} scenarios: {scenario_names}.", + "// Regenerate from an HRX command program dump rather than editing by hand.", + "", + *(f"target.decl @{target}" for target in targets), + "", + *decls, + "", + "", + ] + for scenario in scenarios: + lines.append(f"// Scenario: {scenario.scenario}") + lines.append("") + for case in scenario.cases: + lines.append(case) + lines.append("") + symbol_match = re.match(r"check\.case public @(.+?)_case", case) + if symbol_match: + benchmark_symbol = symbol_match.group(1) + lines.append(f"check.benchmark<@{benchmark_symbol}_case> @{benchmark_symbol}") + lines.append("") + + model_file.parent.mkdir(parents=True, exist_ok=True) + model_file.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8") + + +def write_manifest(manifest_file: Path, model_file: Path, model_slug: str, scenario: ScenarioBenchmarks) -> None: + write_json( + manifest_file, + { + "schema": "ggml-hrx-model-loom-benchmarks-v2", + "model": model_slug, + "scenario": scenario.scenario, + "loom_source": str(model_file.relative_to(HRX_DIR)), + "command_count": scenario.command_count, + "dispatch_count": scenario.dispatch_count, + "generated_count": scenario.generated_count, + "kernel_counts": scenario.kernel_counts, + "dispatches": scenario.dispatches, + }, + ) + + +def parse_scenario_dump(value: str) -> tuple[str, Path]: + scenario, separator, dump_dir = value.partition("=") + if not separator or not scenario or not dump_dir: + fail(f"invalid --scenario-dump value {value!r}; expected =") + return scenario, Path(dump_dir) + + +def scenario_dumps_from_args(args: argparse.Namespace) -> list[tuple[str, Path]]: + if args.scenario_dump: + if args.scenario is not None or args.dump_dir is not None: + fail("--scenario-dump cannot be combined with --scenario or --dump-dir") + if args.output_manifest is not None: + fail("--output-manifest cannot be used with multiple --scenario-dump values") + return [parse_scenario_dump(value) for value in args.scenario_dump] + if args.scenario is None or args.dump_dir is None: + fail("either --scenario and --dump-dir, or one or more --scenario-dump values, are required") + return [(args.scenario, args.dump_dir)] + + +def generate(args: argparse.Namespace) -> None: + exports = load_exports() + model_slug = args.model + model_file = args.output_loom or (BENCHMARK_DIR / "loom" / f"{model_slug}.loom") + scenario_dumps = scenario_dumps_from_args(args) + scenarios: list[ScenarioBenchmarks] = [] + used_exports: dict[str, dict[str, Any]] = {} + for scenario, dump_dir in scenario_dumps: + commands = load_program_commands(dump_dir) + generated = generate_scenario(model_slug, scenario, commands, exports) + scenarios.append(generated) + used_exports.update(generated.used_exports) + + write_model_loom(model_file, model_slug, scenarios, used_exports) + for scenario in scenarios: + manifest_file = args.output_manifest or (BENCHMARK_DIR / "loom" / f"{model_slug}.{scenario.scenario}.json") + write_manifest(manifest_file, model_file, model_slug, scenario) + print(f"wrote {manifest_file}") + + print(f"wrote {model_file}") + for scenario in scenarios: + print( + f"{scenario.scenario}: commands={scenario.command_count} " + f"dispatches={scenario.dispatch_count} generated={scenario.generated_count}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Model slug, for example llama32_3b_f16.") + parser.add_argument("--scenario", help="Scenario slug, for example pp512.") + parser.add_argument("--dump-dir", type=Path, help="Directory containing HRX program-*/program.json dumps.") + parser.add_argument( + "--scenario-dump", + action="append", + help="Scenario and dump directory as =. May be repeated.", + ) + parser.add_argument("--output-loom", type=Path, help="Generated model Loom file.") + parser.add_argument("--output-manifest", type=Path, help="Generated scenario manifest.") + generate(parser.parse_args()) + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 000000000000..9f14062192b4 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl.work.loom @@ -0,0 +1,1955 @@ +// Generated by tools/benchmarks/generate-model-benchmarks.py for qwen38_27b_udq4kxl scenarios: tg_c5. +// Regenerate from an HRX command program dump rather than editing by hand. + +target.decl @ggml_binary_f32_gfx11_wave64 +target.decl @ggml_binary_swiglu_i4_gfx11_wave32 +target.decl @ggml_copy_f32_gfx11_wave64 +target.decl @ggml_flash_attention_gfx11_wave64 +target.decl @ggml_get_rows_f32_gfx11_wave64 +target.decl @ggml_mul_mat_f32_f32_decode_gfx11_wave64 +target.decl @ggml_rmsnorm_binary_gfx11_wave32 +target.decl @ggml_rmsnorm_gfx11_wave32 +target.decl @ggml_scale_bias_f32_gfx11_wave64 +target.decl @ggml_set_rows_gfx11_wave64 +target.decl @qwen3_moe_dense_gfx11_wave64 + +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_binary_f32_gfx11_wave64) @ggml_binary_f32(%element_count: index) launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_binary_swiglu_i4_gfx11_wave32) @ggml_binary_swiglu_symmetric_i4_k32() launch(%lhs: buffer, %rhs: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_concat_dim0_f32() launch(%lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_copy_f32(%element_count: index) launch(%element_count: index, %source: buffer, %output: buffer) +kernel.decl @ggml_fill_negative_f32(%element_count: index) launch(%element_count: index, %output: buffer) +kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) +kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@ggml_mul_mat_f32_f32_decode_gfx11_wave64) @ggml_mul_mat_f32_f32_decode_wave64(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@qwen3_moe_dense_gfx11_wave64) @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer) +kernel.decl @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_selected_refine_token1(%token_count: index, %candidate_count: index) launch(%token_count: index, %candidate_count: index, %input: buffer, %weight: buffer, %candidates: buffer, %exact_output: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_token1_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_symmetric_i2_scan_token1() launch(%weight: buffer, %output: buffer, %qact: buffer, %scales: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma() launch(%first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_quantize_f32_symmetric_i4_k32() launch(%input: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count: index) launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +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_scale_bias_f32_gfx11_wave64) @ggml_scale_bias_f32(%element_count: index) launch(%element_count: index, %input: 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 @ggml_top_k128_f32_reduce_gather_register(%element_count: index) launch(%element_count: index, %partial_values: buffer, %partial_ids: buffer, %candidate_output: buffer, %value_output: buffer) +kernel.decl @ggml_top_k8_f32_partitions_register(%element_count: index) launch(%element_count: index, %values: buffer, %partial_values: buffer, %partial_ids: buffer) +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(%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 + +check.case public @qwen38_27b_udq4kxl_tg_c5_000_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<8xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<8xi8>, tensor<715161600xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_000_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_000_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<40960xi8>, tensor<81920xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>, tensor<81920xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_003_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %first_weight = check.generate.fill value(0) : tensor<138240xi8> + %second_weight = check.generate.fill value(0) : tensor<138240xi8> + %first_output = check.generate.fill value(0) : tensor<384xi8> + %second_output = check.generate.fill value(0) : tensor<384xi8> + kernel.launch @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma[%token_count](%token_count, %input, %first_weight, %second_weight, %first_output, %second_output) : [index](index, tensor<40960xi8>, tensor<138240xi8>, tensor<138240xi8>, tensor<384xi8>, tensor<384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_003_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_003_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<40960xi8>, tensor<49152xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_005_ggml_scale_bias_f32_case { + %element_count = check.literal value(30720) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_scale_bias_f32[%element_count](%element_count, %input, %output) : [index](index, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_005_ggml_scale_bias_f32_case> @qwen38_27b_udq4kxl_tg_c5_005_ggml_scale_bias_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_006_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(30720) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x30720xf32> + %output = check.generate.fill value(1.0) : tensor<1x30720xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x30720xf32>, tensor<1x30720xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_006_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_006_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<81920xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<81920xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_008_ggml_scale_bias_f32_case { + %element_count = check.literal value(786432) : index + %input = check.generate.fill value(0) : tensor<3145728xi8> + %output = check.generate.fill value(0) : tensor<3145728xi8> + kernel.launch @ggml_scale_bias_f32[%element_count](%element_count, %input, %output) : [index](index, tensor<3145728xi8>, tensor<3145728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_008_ggml_scale_bias_f32_case> @qwen38_27b_udq4kxl_tg_c5_008_ggml_scale_bias_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_009_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(786432) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x786432xf32> + %output = check.generate.fill value(1.0) : tensor<1x786432xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x786432xf32>, tensor<1x786432xf32>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<49152xi8> + %k = check.generate.fill value(0) : tensor<49152xi8> + %v = check.generate.fill value(0) : tensor<65536xi8> + %g = check.generate.fill value(0) : tensor<384xi8> + %beta = check.generate.fill value(0) : tensor<384xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<15728640xi8> + %dst = check.generate.fill value(0) : tensor<49152xi8> + %rms_scales = check.generate.fill value(0) : tensor<384xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<49152xi8>, tensor<49152xi8>, tensor<65536xi8>, tensor<384xi8>, tensor<384xi8>, tensor<3145728xi8>, tensor<15728640xi8>, tensor<49152xi8>, tensor<384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(96) : index + %input = check.generate.fill value(0) : tensor<49152xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<49152xi8>, tensor<512xi8>, tensor<49152xi8>, tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<49152xi8>, tensor<40960xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %residual_output = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<40960xi8>, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<139264xi8> + %second_output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<139264xi8> + %rhs = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<139264xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<139264xi8>, tensor<40960xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>, tensor<40960xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<40960xi8>, tensor<98304xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<40960xi8>, tensor<8192xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_021_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<97280xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<97280xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<49152xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_021_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_021_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_022_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<8192xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<8192xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<8192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_022_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_022_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_023_ggml_set_rows_case { + %token_count = check.literal value(2) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<2x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<2xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<2x1024xf32>, tensor<2xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_023_ggml_set_rows_case> @qwen38_27b_udq4kxl_tg_c5_023_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_tg_c5_024_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(2) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<2x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<2x256xf16> + %output = check.generate.fill value(1.0) : tensor<2x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<2x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<2x256xf16>, tensor<2x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_024_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_024_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_026_ggml_binary_f32_case { + %element_count = check.literal value(10240) : index + %lhs = check.generate.fill value(2.0) : tensor<10240xf32> + %rhs = check.generate.fill value(3.0) : tensor<10240xf32> + %output = check.generate.fill value(0.0) : tensor<10240xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<10240xf32>, tensor<10240xf32>, tensor<10240xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_026_ggml_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_026_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_027_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(2.0) : tensor<2x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<2x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<2x5120xf32>, tensor<5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_027_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_027_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_028_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<2x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_028_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_028_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<20480xi8>, tensor<1042944000xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_030_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<40960xi8>, tensor<40960xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_030_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_tg_c5_030_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_032_ggml_get_rows_f32_case { + %token_count = check.literal value(5) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<20xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<20xi8>, tensor<715161600xi8>, tensor<102400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_032_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_032_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<102400xi8>, tensor<20480xi8>, tensor<102400xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + %partial = check.generate.fill value(0) : tensor<204800xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<102400xi8>, tensor<204800xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>, tensor<204800xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_035_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %first_weight = check.generate.fill value(0) : tensor<138240xi8> + %second_weight = check.generate.fill value(0) : tensor<138240xi8> + %first_output = check.generate.fill value(0) : tensor<960xi8> + %second_output = check.generate.fill value(0) : tensor<960xi8> + kernel.launch @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma[%token_count](%token_count, %input, %first_weight, %second_weight, %first_output, %second_output) : [index](index, tensor<102400xi8>, tensor<138240xi8>, tensor<138240xi8>, tensor<960xi8>, tensor<960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_035_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_035_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<102400xi8>, tensor<122880xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<204800xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<204800xi8>, tensor<163840xi8>, tensor<204800xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<172032xi8> + %k = check.generate.fill value(0) : tensor<172032xi8> + %v = check.generate.fill value(0) : tensor<188416xi8> + %g = check.generate.fill value(0) : tensor<960xi8> + %beta = check.generate.fill value(0) : tensor<960xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<53477376xi8> + %dst = check.generate.fill value(0) : tensor<122880xi8> + %rms_scales = check.generate.fill value(0) : tensor<960xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<172032xi8>, tensor<172032xi8>, tensor<188416xi8>, tensor<960xi8>, tensor<960xi8>, tensor<3145728xi8>, tensor<53477376xi8>, tensor<122880xi8>, tensor<960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(240) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<122880xi8>, tensor<512xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<122880xi8>, tensor<102400xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(5) : index + %lhs = check.generate.fill value(0) : tensor<102400xi8> + %rhs = check.generate.fill value(0) : tensor<102400xi8> + %residual_output = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<102400xi8>, tensor<102400xi8>, tensor<102400xi8>, tensor<20480xi8>, tensor<102400xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<348160xi8> + %second_output = check.generate.fill value(0) : tensor<348160xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<348160xi8>, tensor<348160xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<348160xi8> + %rhs = check.generate.fill value(0) : tensor<348160xi8> + %output = check.generate.fill value(0) : tensor<348160xi8> + %quantized_values = check.generate.fill value(0) : tensor<43520xi8> + %scales = check.generate.fill value(0) : tensor<10880xi8> + %sums = check.generate.fill value(0) : tensor<10880xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<348160xi8>, tensor<348160xi8>, tensor<348160xi8>, tensor<43520xi8>, tensor<10880xi8>, tensor<10880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<348160xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<43520xi8> + %scales = check.generate.fill value(0) : tensor<10880xi8> + %sums = check.generate.fill value(0) : tensor<10880xi8> + %partial = check.generate.fill value(0) : tensor<102400xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<348160xi8>, tensor<102400xi8>, tensor<43520xi8>, tensor<10880xi8>, tensor<10880xi8>, tensor<102400xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<245760xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<102400xi8>, tensor<245760xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<102400xi8>, tensor<5611520xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<102400xi8>, tensor<20480xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_049_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<244736xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<80xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<244736xi8>, tensor<1024xi8>, tensor<80xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_049_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_049_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_050_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<80xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<20480xi8>, tensor<1024xi8>, tensor<80xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_050_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_050_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_051_ggml_set_rows_case { + %token_count = check.literal value(5) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<5x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<5xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<5x1024xf32>, tensor<5xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_051_ggml_set_rows_case> @qwen38_27b_udq4kxl_tg_c5_051_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_tg_c5_052_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(5) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<5x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<5x256xf16> + %output = check.generate.fill value(1.0) : tensor<5x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<5x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<5x256xf16>, tensor<5x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_052_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_052_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<122880xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_054_ggml_binary_f32_case { + %element_count = check.literal value(25600) : index + %lhs = check.generate.fill value(2.0) : tensor<25600xf32> + %rhs = check.generate.fill value(3.0) : tensor<25600xf32> + %output = check.generate.fill value(0.0) : tensor<25600xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<25600xf32>, tensor<25600xf32>, tensor<25600xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_054_ggml_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_054_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_055_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(2.0) : tensor<5x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<5x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<5x5120xf32>, tensor<5120xf32>, tensor<5x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_055_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_055_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_056_ggml_get_rows_f32_case { + %token_count = check.literal value(5) : index + %row_count = check.literal value(5) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<5xi32> + %weight = check.generate.fill value(0.0) : tensor<5x5120xf32> + %output = check.generate.fill value(1.0) : tensor<5x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<5xi32>, tensor<5x5120xf32>, tensor<5x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_056_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_056_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_058_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<102400xi8> + %rhs = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<102400xi8>, tensor<102400xi8>, tensor<204800xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_058_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_tg_c5_058_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_059_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<204800xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<204800xi8>, tensor<56115200xi8>, tensor<102400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_059_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_059_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_060_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<2xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<2x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<2xi32>, tensor<2x5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_060_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_060_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_062_ggml_get_rows_f32_case { + %token_count = check.literal value(16) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<64xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<64xi8>, tensor<715161600xi8>, tensor<327680xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_062_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_062_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_063_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<327680xi8>, tensor<20480xi8>, tensor<327680xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_063_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_063_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_064_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + %partial = check.generate.fill value(0) : tensor<655360xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<327680xi8>, tensor<655360xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>, tensor<655360xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_064_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_064_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_065_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %first_weight = check.generate.fill value(0) : tensor<138240xi8> + %second_weight = check.generate.fill value(0) : tensor<138240xi8> + %first_output = check.generate.fill value(0) : tensor<3072xi8> + %second_output = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma[%token_count](%token_count, %input, %first_weight, %second_weight, %first_output, %second_output) : [index](index, tensor<327680xi8>, tensor<138240xi8>, tensor<138240xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_065_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_065_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_066_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<327680xi8>, tensor<393216xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_066_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_066_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_067_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<655360xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<655360xi8>, tensor<163840xi8>, tensor<655360xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_068_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_tg_c5_068_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_069_llm_gated_delta_net_f32_wmma_head128_case { + %q = check.generate.fill value(0) : tensor<417792xi8> + %k = check.generate.fill value(0) : tensor<417792xi8> + %v = check.generate.fill value(0) : tensor<434176xi8> + %g = check.generate.fill value(0) : tensor<2112xi8> + %beta = check.generate.fill value(0) : tensor<2112xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<3416064xi8> + %rms_scales = check.generate.fill value(0) : tensor<2112xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128[](%q, %k, %v, %g, %beta, %state_in, %dst, %rms_scales) : [](tensor<417792xi8>, tensor<417792xi8>, tensor<434176xi8>, tensor<2112xi8>, tensor<2112xi8>, tensor<3145728xi8>, tensor<3416064xi8>, tensor<2112xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_069_llm_gated_delta_net_f32_wmma_head128_case> @qwen38_27b_udq4kxl_tg_c5_069_llm_gated_delta_net_f32_wmma_head128 + +check.case public @qwen38_27b_udq4kxl_tg_c5_070_ggml_copy_f32_case { + %element_count = check.literal value(786432) : index + %source = check.generate.fill value(0) : tensor<3145728xi8> + %output = check.generate.fill value(0) : tensor<3145728xi8> + kernel.launch @ggml_copy_f32[%element_count](%element_count, %source, %output) : [index](index, tensor<3145728xi8>, tensor<3145728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_070_ggml_copy_f32_case> @qwen38_27b_udq4kxl_tg_c5_070_ggml_copy_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_071_llm_gated_delta_net_f32_wmma_head128_inplace_case { + %q = check.generate.fill value(0) : tensor<8192xi8> + %k = check.generate.fill value(0) : tensor<8192xi8> + %v = check.generate.fill value(0) : tensor<24576xi8> + %g = check.generate.fill value(0) : tensor<192xi8> + %beta = check.generate.fill value(0) : tensor<192xi8> + %state_inout = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<24576xi8> + %rms_scales = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_inplace[](%q, %k, %v, %g, %beta, %state_inout, %dst, %rms_scales) : [](tensor<8192xi8>, tensor<8192xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3145728xi8>, tensor<24576xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_071_llm_gated_delta_net_f32_wmma_head128_inplace_case> @qwen38_27b_udq4kxl_tg_c5_071_llm_gated_delta_net_f32_wmma_head128_inplace + +check.case public @qwen38_27b_udq4kxl_tg_c5_072_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(768) : index + %input = check.generate.fill value(0) : tensor<393216xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<393216xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<393216xi8>, tensor<512xi8>, tensor<393216xi8>, tensor<393216xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_072_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_072_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_073_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<393216xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<393216xi8>, tensor<327680xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_073_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_073_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_074_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(16) : index + %lhs = check.generate.fill value(0) : tensor<327680xi8> + %rhs = check.generate.fill value(0) : tensor<327680xi8> + %residual_output = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<327680xi8>, tensor<327680xi8>, tensor<327680xi8>, tensor<20480xi8>, tensor<327680xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_074_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_074_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_075_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<1114112xi8> + %second_output = check.generate.fill value(0) : tensor<1114112xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<1114112xi8>, tensor<1114112xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_075_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_075_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_076_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<1114112xi8> + %rhs = check.generate.fill value(0) : tensor<1114112xi8> + %output = check.generate.fill value(0) : tensor<1114112xi8> + %quantized_values = check.generate.fill value(0) : tensor<139264xi8> + %scales = check.generate.fill value(0) : tensor<34816xi8> + %sums = check.generate.fill value(0) : tensor<34816xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<1114112xi8>, tensor<1114112xi8>, tensor<1114112xi8>, tensor<139264xi8>, tensor<34816xi8>, tensor<34816xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_076_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_076_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_077_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<1114112xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<139264xi8> + %scales = check.generate.fill value(0) : tensor<34816xi8> + %sums = check.generate.fill value(0) : tensor<34816xi8> + %partial = check.generate.fill value(0) : tensor<327680xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<1114112xi8>, tensor<327680xi8>, tensor<139264xi8>, tensor<34816xi8>, tensor<34816xi8>, tensor<327680xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_077_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_077_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_078_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<786432xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<327680xi8>, tensor<786432xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_078_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_078_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_079_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<327680xi8>, tensor<5611520xi8>, tensor<65536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_079_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_079_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_080_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<327680xi8>, tensor<65536xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_080_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_080_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_081_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<785408xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<785408xi8>, tensor<1024xi8>, tensor<256xi8>, tensor<393216xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_081_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_081_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_082_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<65536xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<65536xi8>, tensor<1024xi8>, tensor<256xi8>, tensor<65536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_082_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_082_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_083_ggml_set_rows_case { + %token_count = check.literal value(16) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<16x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<16xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<16x1024xf32>, tensor<16xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_083_ggml_set_rows_case> @qwen38_27b_udq4kxl_tg_c5_083_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_tg_c5_084_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(16) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<16x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<16x256xf16> + %output = check.generate.fill value(1.0) : tensor<16x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<16x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<16x256xf16>, tensor<16x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_084_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_084_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_085_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<393216xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_085_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_085_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_086_ggml_binary_f32_case { + %element_count = check.literal value(81920) : index + %lhs = check.generate.fill value(2.0) : tensor<81920xf32> + %rhs = check.generate.fill value(3.0) : tensor<81920xf32> + %output = check.generate.fill value(0.0) : tensor<81920xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<81920xf32>, tensor<81920xf32>, tensor<81920xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_086_ggml_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_086_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_087_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(2.0) : tensor<16x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<16x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<16x5120xf32>, tensor<5120xf32>, tensor<16x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_087_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_087_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_088_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<327680xi8> + %rhs = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<327680xi8>, tensor<327680xi8>, tensor<655360xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_088_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_tg_c5_088_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_089_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<655360xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<655360xi8>, tensor<56115200xi8>, tensor<327680xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_089_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_089_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_090_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<16xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<16xi8>, tensor<715161600xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_090_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_090_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_091_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_091_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_091_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_092_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + %partial = check.generate.fill value(0) : tensor<163840xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>, tensor<163840xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_092_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_092_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_093_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %first_weight = check.generate.fill value(0) : tensor<138240xi8> + %second_weight = check.generate.fill value(0) : tensor<138240xi8> + %first_output = check.generate.fill value(0) : tensor<768xi8> + %second_output = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma[%token_count](%token_count, %input, %first_weight, %second_weight, %first_output, %second_output) : [index](index, tensor<81920xi8>, tensor<138240xi8>, tensor<138240xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_093_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_093_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_094_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<81920xi8>, tensor<98304xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_094_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_094_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_095_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<163840xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_096_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_tg_c5_096_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_097_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<131072xi8> + %k = check.generate.fill value(0) : tensor<131072xi8> + %v = check.generate.fill value(0) : tensor<147456xi8> + %g = check.generate.fill value(0) : tensor<768xi8> + %beta = check.generate.fill value(0) : tensor<768xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<40894464xi8> + %dst = check.generate.fill value(0) : tensor<98304xi8> + %rms_scales = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<131072xi8>, tensor<131072xi8>, tensor<147456xi8>, tensor<768xi8>, tensor<768xi8>, tensor<3145728xi8>, tensor<40894464xi8>, tensor<98304xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_097_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_tg_c5_097_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_tg_c5_098_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(192) : index + %input = check.generate.fill value(0) : tensor<98304xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<98304xi8>, tensor<512xi8>, tensor<98304xi8>, tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_098_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_098_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<98304xi8>, tensor<81920xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_100_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %residual_output = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<81920xi8>, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_100_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_100_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_101_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<278528xi8> + %second_output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_101_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_101_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_102_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<278528xi8> + %rhs = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<278528xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_102_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_102_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_103_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<278528xi8>, tensor<81920xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>, tensor<81920xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_103_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_103_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_104_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<196608xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<81920xi8>, tensor<196608xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_104_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_104_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_105_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<5611520xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_105_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_105_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_106_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<81920xi8>, tensor<16384xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_106_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_106_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_107_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<195584xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<195584xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<98304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_107_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_107_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_108_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<16384xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<16384xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_108_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_108_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_109_ggml_set_rows_case { + %token_count = check.literal value(4) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<4x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<4xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<4x1024xf32>, tensor<4xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_109_ggml_set_rows_case> @qwen38_27b_udq4kxl_tg_c5_109_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_tg_c5_110_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<4x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<4x256xf16> + %output = check.generate.fill value(1.0) : tensor<4x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<4x256xf16>, tensor<4x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_110_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_110_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_111_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_111_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_111_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_112_ggml_binary_f32_case { + %element_count = check.literal value(20480) : index + %lhs = check.generate.fill value(2.0) : tensor<20480xf32> + %rhs = check.generate.fill value(3.0) : tensor<20480xf32> + %output = check.generate.fill value(0.0) : tensor<20480xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<20480xf32>, tensor<20480xf32>, tensor<20480xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_112_ggml_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_112_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_113_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(2.0) : tensor<4x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<4x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<4x5120xf32>, tensor<5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_113_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_113_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_114_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<4x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_114_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_114_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_115_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<81920xi8>, tensor<81920xi8>, tensor<163840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_115_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_tg_c5_115_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_116_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<163840xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<163840xi8>, tensor<56115200xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_116_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_116_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_117_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<4xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi8>, tensor<715161600xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_117_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_117_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_118_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<1x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<1x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<1x5120xf32>, tensor<5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_118_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_118_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_119_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<20480xi8> + %rhs = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<20480xi8>, tensor<20480xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_119_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_tg_c5_119_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_120_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(10240) : index + %output_size = check.literal value(5120) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<43008000xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<40960xi8>, tensor<43008000xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_120_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_tg_c5_120_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_tg_c5_121_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_121_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_121_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_122_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<20480xi8>, tensor<49152xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_122_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_122_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_123_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(5120) : index + %output_size = check.literal value(1024) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<4300800xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<20480xi8>, tensor<4300800xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_123_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_tg_c5_123_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_tg_c5_124_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<20480xi8>, tensor<4096xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_124_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_124_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_125_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<48128xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<16xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<48128xi8>, tensor<1024xi8>, tensor<16xi8>, tensor<24576xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_125_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_125_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_126_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<4096xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<16xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<4096xi8>, tensor<1024xi8>, tensor<16xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_126_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_tg_c5_126_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_127_ggml_set_rows_case { + %token_count = check.literal value(1) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<1x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<1xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<1x1024xf32>, tensor<1xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_127_ggml_set_rows_case> @qwen38_27b_udq4kxl_tg_c5_127_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_tg_c5_128_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<1x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<1x256xf16> + %output = check.generate.fill value(1.0) : tensor<1x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<1x256xf16>, tensor<1x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_128_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_tg_c5_128_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_129_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<24576xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_129_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_129_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<24576xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<24576xi8>, tensor<20480xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_131_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(1) : index + %lhs = check.generate.fill value(0) : tensor<20480xi8> + %rhs = check.generate.fill value(0) : tensor<20480xi8> + %residual_output = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_131_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_131_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<69632xi8> + %second_output = check.generate.fill value(0) : tensor<69632xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<69632xi8>, tensor<69632xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_133_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<69632xi8> + %rhs = check.generate.fill value(0) : tensor<69632xi8> + %output = check.generate.fill value(0) : tensor<69632xi8> + %quantized_values = check.generate.fill value(0) : tensor<8704xi8> + %scales = check.generate.fill value(0) : tensor<2176xi8> + %sums = check.generate.fill value(0) : tensor<2176xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<69632xi8>, tensor<69632xi8>, tensor<69632xi8>, tensor<8704xi8>, tensor<2176xi8>, tensor<2176xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_133_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_133_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_134_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<69632xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<8704xi8> + %scales = check.generate.fill value(0) : tensor<2176xi8> + %sums = check.generate.fill value(0) : tensor<2176xi8> + %partial = check.generate.fill value(0) : tensor<20480xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<69632xi8>, tensor<20480xi8>, tensor<8704xi8>, tensor<2176xi8>, tensor<2176xi8>, tensor<20480xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_134_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_134_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_135_ggml_binary_f32_case { + %element_count = check.literal value(5120) : index + %lhs = check.generate.fill value(2.0) : tensor<5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<5120xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<5120xf32>, tensor<5120xf32>, tensor<5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_135_ggml_binary_f32_case> @qwen38_27b_udq4kxl_tg_c5_135_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_136_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(1) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<1x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<1x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_136_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_136_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_138_ggml_mul_mat_q6_k_symmetric_i2_scan_token1_case { + %weight = check.generate.fill value(0) : tensor<1380659200xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + %qact = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_q6_k_symmetric_i2_scan_token1[](%weight, %output, %qact, %scales) : [](tensor<1380659200xi8>, tensor<993280xi8>, tensor<2560xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_138_ggml_mul_mat_q6_k_symmetric_i2_scan_token1_case> @qwen38_27b_udq4kxl_tg_c5_138_ggml_mul_mat_q6_k_symmetric_i2_scan_token1 + +check.case public @qwen38_27b_udq4kxl_tg_c5_139_ggml_top_k8_f32_partitions_register_case { + %element_count = check.literal value(248320) : index + %values = check.generate.fill value(0) : tensor<993280xi8> + %partial_values = check.generate.fill value(0) : tensor<4096xi8> + %partial_ids = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_top_k8_f32_partitions_register[%element_count](%element_count, %values, %partial_values, %partial_ids) : [index](index, tensor<993280xi8>, tensor<4096xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_139_ggml_top_k8_f32_partitions_register_case> @qwen38_27b_udq4kxl_tg_c5_139_ggml_top_k8_f32_partitions_register + +check.case public @qwen38_27b_udq4kxl_tg_c5_140_ggml_top_k128_f32_reduce_gather_register_case { + %element_count = check.literal value(248320) : index + %partial_values = check.generate.fill value(0) : tensor<4096xi8> + %partial_ids = check.generate.fill value(0) : tensor<4096xi8> + %candidate_output = check.generate.fill value(0) : tensor<512xi8> + %value_output = check.generate.fill value(0) : tensor<512xi8> + kernel.launch @ggml_top_k128_f32_reduce_gather_register[%element_count](%element_count, %partial_values, %partial_ids, %candidate_output, %value_output) : [index](index, tensor<4096xi8>, tensor<4096xi8>, tensor<512xi8>, tensor<512xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_140_ggml_top_k128_f32_reduce_gather_register_case> @qwen38_27b_udq4kxl_tg_c5_140_ggml_top_k128_f32_reduce_gather_register + +check.case public @qwen38_27b_udq4kxl_tg_c5_141_ggml_fill_negative_f32_case { + %element_count = check.literal value(248320) : index + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_fill_negative_f32[%element_count](%element_count, %output) : [index](index, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_141_ggml_fill_negative_f32_case> @qwen38_27b_udq4kxl_tg_c5_141_ggml_fill_negative_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_142_ggml_mul_mat_q6_k_packed_selected_refine_token1_case { + %token_count = check.literal value(1) : index + %candidate_count = check.literal value(64) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1380659200xi8> + %candidates = check.generate.fill value(0) : tensor<256xi8> + %exact_output = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_selected_refine_token1[%token_count, %candidate_count](%token_count, %candidate_count, %input, %weight, %candidates, %exact_output, %output) : [index, index](index, index, tensor<20480xi8>, tensor<1380659200xi8>, tensor<256xi8>, tensor<256xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_142_ggml_mul_mat_q6_k_packed_selected_refine_token1_case> @qwen38_27b_udq4kxl_tg_c5_142_ggml_mul_mat_q6_k_packed_selected_refine_token1 + +check.case public @qwen38_27b_udq4kxl_tg_c5_143_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>, tensor<40960xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_143_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_tg_c5_143_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_144_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(5120) : index + %output_size = check.literal value(48) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<138240xi8> + %output = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<20480xi8>, tensor<138240xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_144_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_tg_c5_144_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_tg_c5_145_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<20480xi8>, tensor<24576xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_145_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_tg_c5_145_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_tg_c5_146_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<40960xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<40960xi8>, tensor<163840xi8>, tensor<40960xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_147_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_tg_c5_147_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_148_llm_gated_delta_net_f32_wmma_head128_inplace_case { + %q = check.generate.fill value(0) : tensor<8192xi8> + %k = check.generate.fill value(0) : tensor<8192xi8> + %v = check.generate.fill value(0) : tensor<24576xi8> + %g = check.generate.fill value(0) : tensor<192xi8> + %beta = check.generate.fill value(0) : tensor<192xi8> + %state_inout = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<24576xi8> + %rms_scales = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_inplace[](%q, %k, %v, %g, %beta, %state_inout, %dst, %rms_scales) : [](tensor<8192xi8>, tensor<8192xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3145728xi8>, tensor<24576xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_148_llm_gated_delta_net_f32_wmma_head128_inplace_case> @qwen38_27b_udq4kxl_tg_c5_148_llm_gated_delta_net_f32_wmma_head128_inplace + +check.case public @qwen38_27b_udq4kxl_tg_c5_149_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(48) : index + %input = check.generate.fill value(0) : tensor<24576xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<24576xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<24576xi8>, tensor<512xi8>, tensor<24576xi8>, tensor<24576xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_149_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_tg_c5_149_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_tg_c5_150_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<4xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<4x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi32>, tensor<4x5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_tg_c5_150_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_tg_c5_150_ggml_get_rows_f32 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 new file mode 100644 index 000000000000..df898f2bde8d --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/qwen38_27b_udq4kxl_recurrent.work.loom @@ -0,0 +1,390 @@ +// Generated by tools/benchmarks/generate-model-benchmarks.py for qwen38_27b_udq4kxl_recurrent scenarios: tg_c5. +// Regenerate from an HRX command program dump rather than editing by hand. + +target.decl @ggml_binary_f32_gfx11_wave64 +target.decl @ggml_binary_swiglu_i4_gfx11_wave32 +target.decl @ggml_copy_f32_gfx11_wave64 +target.decl @ggml_flash_attention_gfx11_wave64 +target.decl @ggml_get_rows_f32_gfx11_wave64 +target.decl @ggml_rmsnorm_binary_gfx11_wave32 +target.decl @ggml_rmsnorm_gfx11_wave32 +target.decl @ggml_set_rows_gfx11_wave64 +target.decl @qwen3_moe_dense_gfx11_wave64 + +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_binary_f32_gfx11_wave64) @ggml_binary_f32(%element_count: index) launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_binary_swiglu_i4_gfx11_wave32) @ggml_binary_swiglu_symmetric_i4_k32() launch(%lhs: buffer, %rhs: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_concat_dim0_f32() launch(%lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) +kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@qwen3_moe_dense_gfx11_wave64) @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma() launch(%first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_quantize_f32_symmetric_i4_k32() launch(%input: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count: index) launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +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(%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 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_000_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<8xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<8xi8>, tensor<715161600xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_000_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_000_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<40960xi8>, tensor<81920xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>, tensor<81920xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_003_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %first_weight = check.generate.fill value(0) : tensor<138240xi8> + %second_weight = check.generate.fill value(0) : tensor<138240xi8> + %first_output = check.generate.fill value(0) : tensor<384xi8> + %second_output = check.generate.fill value(0) : tensor<384xi8> + kernel.launch @ggml_mul_mat_q4_k_f32_adjacent_dual_wmma[%token_count](%token_count, %input, %first_weight, %second_weight, %first_output, %second_output) : [index](index, tensor<40960xi8>, tensor<138240xi8>, tensor<138240xi8>, tensor<384xi8>, tensor<384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_003_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_003_ggml_mul_mat_q4_k_f32_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<40960xi8>, tensor<49152xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_005_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(30720) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x30720xf32> + %output = check.generate.fill value(1.0) : tensor<1x30720xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x30720xf32>, tensor<1x30720xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_005_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_005_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_006_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<81920xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<81920xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_006_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_006_llm_ssm_conv_dconv4_silu_rollback_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_007_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(786432) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x786432xf32> + %output = check.generate.fill value(1.0) : tensor<1x786432xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x786432xf32>, tensor<1x786432xf32>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_008_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_008_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_009_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<49152xi8> + %k = check.generate.fill value(0) : tensor<49152xi8> + %v = check.generate.fill value(0) : tensor<65536xi8> + %g = check.generate.fill value(0) : tensor<384xi8> + %beta = check.generate.fill value(0) : tensor<384xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<15728640xi8> + %dst = check.generate.fill value(0) : tensor<49152xi8> + %rms_scales = check.generate.fill value(0) : tensor<384xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<49152xi8>, tensor<49152xi8>, tensor<65536xi8>, tensor<384xi8>, tensor<384xi8>, tensor<3145728xi8>, tensor<15728640xi8>, tensor<49152xi8>, tensor<384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_009_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_009_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(96) : index + %input = check.generate.fill value(0) : tensor<49152xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<49152xi8>, tensor<512xi8>, tensor<49152xi8>, tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_011_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<49152xi8>, tensor<40960xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_011_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_011_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_012_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %residual_output = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<40960xi8>, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_012_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_012_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<139264xi8> + %second_output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_014_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<139264xi8> + %rhs = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<139264xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_014_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_014_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<139264xi8>, tensor<40960xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>, tensor<40960xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_016_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<40960xi8>, tensor<98304xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_016_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_016_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<40960xi8>, tensor<8192xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_019_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<97280xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<97280xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<49152xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_019_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_019_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_020_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<8192xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<8192xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<8192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_020_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_020_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_021_ggml_set_rows_case { + %token_count = check.literal value(2) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<2x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<2xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<2x1024xf32>, tensor<2xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_021_ggml_set_rows_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_021_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_022_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(2) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<2x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<2x256xf16> + %output = check.generate.fill value(1.0) : tensor<2x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<2x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<2x256xf16>, tensor<2x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_022_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_022_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_023_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_023_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_023_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_024_ggml_binary_f32_case { + %element_count = check.literal value(10240) : index + %lhs = check.generate.fill value(2.0) : tensor<10240xf32> + %rhs = check.generate.fill value(3.0) : tensor<10240xf32> + %output = check.generate.fill value(0.0) : tensor<10240xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<10240xf32>, tensor<10240xf32>, tensor<10240xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_024_ggml_binary_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_024_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_025_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(2.0) : tensor<2x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<2x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<2x5120xf32>, tensor<5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_025_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_025_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_026_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<2xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<2x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<2xi32>, tensor<2x5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_026_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_026_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_recurrent_tg_c5_028_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<40960xi8>, tensor<40960xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_recurrent_tg_c5_028_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_recurrent_tg_c5_028_ggml_concat_dim0_f32 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 new file mode 100644 index 000000000000..d6f812883b17 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.json @@ -0,0 +1,5792 @@ +{ + "command_count": 8673, + "dispatch_count": 181, + "dispatches": [ + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_000_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "2", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 2 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_001_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_004_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_005_ggml_scale_bias_f32", + "compile_parameters": { + "ggml.scale.bias": "0", + "ggml.scale.scale": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "element_count": 30720 + }, + "kernel": "loom_libs:ggml_scale_bias_f32", + "library_sources": [], + "primary_sources": [ + "ops/scale_bias_f32.loom" + ], + "sources": [ + "ops/scale_bias_f32.loom" + ], + "symbol": "ggml_scale_bias_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_006_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "30720", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 384, + "integer_parameters": { + "hidden_size": 30720, + "row_count": 20, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_007_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "2", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_008_ggml_scale_bias_f32", + "compile_parameters": { + "ggml.scale.bias": "0", + "ggml.scale.scale": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "element_count": 786432 + }, + "kernel": "loom_libs:ggml_scale_bias_f32", + "library_sources": [], + "primary_sources": [ + "ops/scale_bias_f32.loom" + ], + "sources": [ + "ops/scale_bias_f32.loom" + ], + "symbol": "ggml_scale_bias_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_009_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "786432", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 384, + "integer_parameters": { + "hidden_size": 786432, + "row_count": 20, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_010_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 96 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_011_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "20480", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "96", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "2", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "20480", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "token_count": 96 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_013_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_014_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 256, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_016_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_018_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_019_ggml_mul_mat_f32_f32_wmma", + "compile_parameters": { + "ggml.mul_mat.input_size": "5120", + "ggml.mul_mat.output_accumulation": "0", + "ggml.mul_mat.output_size": "1024", + "ggml.mul_mat.output_unary_op": "23", + "ggml.mul_mat.weight_format": "6", + "ggml.workload.token_capacity": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_wmma.loom", + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_020_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_021_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "24576", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "2", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "12288", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_022_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "2048", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "2", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "2048", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_023_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 68, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 2 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_024_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 2 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_025_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_026_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "element_count": 10240 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_027_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 8, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_028_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 2, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_030_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_031_ggml_mul_mat_f32_f32_wmma", + "compile_parameters": { + "ggml.mul_mat.input_size": "10240", + "ggml.mul_mat.output_accumulation": "0", + "ggml.mul_mat.output_size": "5120", + "ggml.mul_mat.output_unary_op": "23", + "ggml.mul_mat.weight_format": "6", + "ggml.workload.token_capacity": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_wmma.loom", + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_032_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "4", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_033_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_036_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_037_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "4", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_038_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 192 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_039_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "40960", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "192", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "4", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "40960", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "token_count": 192 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_041_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_042_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 255, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_044_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_046_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_048_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_049_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "49152", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "4", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "24576", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_050_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "4096", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "4", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "4096", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_051_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_052_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 4 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_053_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_054_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "element_count": 20480 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_055_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 5, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_056_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "4", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 4, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_058_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "5", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 5 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_059_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_062_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_063_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "5", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_064_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 240 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_065_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "51200", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "240", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "5", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "51200", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 240 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_067_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_068_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 128, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_070_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_072_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_074_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_075_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "61440", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "5", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "30720", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_076_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "5120", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "5", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "5120", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_077_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 5 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_078_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 5 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_079_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_080_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "element_count": 25600 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_081_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_082_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "5", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 5, + "token_count": 5 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_084_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_086_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "3", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 3 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_087_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_090_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_091_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "3", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_092_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 144 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_093_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "30720", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "144", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "3", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "30720", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 144 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_095_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_096_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 128, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_098_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_100_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_102_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_103_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "36864", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "3", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "18432", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_104_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "3072", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "3", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "3072", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_105_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 3 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_106_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 3 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_107_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_108_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "element_count": 15360 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_109_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_110_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "3", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 3, + "token_count": 3 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_112_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_114_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "16", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 16 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_115_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_118_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_119_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "16", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_120_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 768 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_121_llm_gated_delta_net_f32_wmma_head128", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "163840", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "768", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.token_count": "11", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "163840", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_122_ggml_copy_f32", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 288, + "integer_parameters": { + "element_count": 786432 + }, + "kernel": "loom_libs:ggml_copy_f32", + "library_sources": [ + "motifs/copy/vector_load_store.loom" + ], + "primary_sources": [ + "ops/copy/copy_f32.loom" + ], + "sources": [ + "ops/copy/copy_f32.loom", + "motifs/copy/vector_load_store.loom" + ], + "symbol": "ggml_copy_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_123_llm_gated_delta_net_f32_wmma_head128_inplace", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "163840", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "768", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.token_count": "1", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "163840", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 240, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_inplace", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_inplace", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 768 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_125_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_126_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 128, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_128_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_130_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_132_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_133_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "196608", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "16", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "98304", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_134_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "16384", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "16", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "16384", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_135_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 16 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_136_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 16 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_137_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_138_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "element_count": 81920 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_139_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_140_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_142_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 4, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_143_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_145_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_146_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 7, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_147_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_148_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "5120", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "6" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "input_size": 10240, + "output_size": 5120, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_149_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_150_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_151_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "1024", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "6" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": { + "input_size": 5120, + "output_size": 1024, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_152_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_153_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "12288", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "1", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "6144", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_154_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "1024", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "1", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "1024", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_155_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 36, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_156_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 1 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_157_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_158_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_159_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_161_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_163_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "element_count": 5120 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_164_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 1, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_165_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "5120", + "ggml.quantize_symmetric_i4_k32.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_166_ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "compile_parameters": { + "ggml.mul_mat_q6_k_shortlist.input_size": "5120", + "ggml.mul_mat_q6_k_shortlist.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_167_ggml_top_k8_f32_partitions_register", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 248320 + }, + "kernel": "loom_libs:ggml_top_k8_f32_partitions_register", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_top_k8_f32_partitions_register", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_168_ggml_top_k128_f32_reduce_gather_register", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 248320 + }, + "kernel": "loom_libs:ggml_top_k128_f32_reduce_gather_register", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_top_k128_f32_reduce_gather_register", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_169_ggml_fill_negative_f32", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 248320 + }, + "kernel": "loom_libs:ggml_fill_negative_f32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_fill_negative_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_170_ggml_mul_mat_q6_k_packed_selected_refine_token1", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320", + "ggml.mul_mat_q6_k_packed.weight_offset": "337715200" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "candidate_count": 64, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_selected_refine_token1", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_selected_refine_token1", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "candidate_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_171_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_172_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "48", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "input_size": 5120, + "output_size": 48, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_173_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_174_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "1", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_175_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 48 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_176_llm_gated_delta_net_f32_wmma_head128_inplace", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "10240", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "48", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.token_count": "1", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "10240", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_inplace", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_inplace", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_177_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 48 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_178_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "2", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 2, + "token_count": 2 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_draft_tg_c1_179_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + } + ], + "generated_count": 180, + "kernel_counts": { + "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32": 1024, + "loom_libs:ggml_binary_f32": 16, + "loom_libs:ggml_binary_swiglu_symmetric_i4_k32": 520, + "loom_libs:ggml_concat_dim0_f32": 8, + "loom_libs:ggml_copy_f32": 288, + "loom_libs:ggml_fill_negative_f32": 1, + "loom_libs:ggml_flash_attention_f32_f16_wmma": 136, + "loom_libs:ggml_get_rows_f32": 793, + "loom_libs:ggml_mul_mat_f32_f32_decode_wave64": 116, + "loom_libs:ggml_mul_mat_f32_f32_wmma": 36, + "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma": 88, + "loom_libs:ggml_mul_mat_q6_k_packed_selected_refine_token1": 2, + "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma": 8, + "loom_libs:ggml_mul_mat_q6_k_symmetric_i2_scan_token1": 1, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma": 856, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma": 904, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma": 1176, + "loom_libs:ggml_quantize_f32_symmetric_i4_k32": 137, + "loom_libs:ggml_rmsnorm_binary_f32": 32, + "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32": 16, + "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32": 384, + "loom_libs:ggml_rmsnorm_mul_rope_f32": 272, + "loom_libs:ggml_scale_bias_f32": 192, + "loom_libs:ggml_select_symmetric_i4_k32_groups": 1, + "loom_libs:ggml_set_rows": 272, + "loom_libs:ggml_top_k128_f32_reduce_gather_register": 1, + "loom_libs:ggml_top_k8_f32_partitions_register": 1, + "loom_libs:llm_gated_delta_net_f32_wmma_head128": 48, + "loom_libs:llm_gated_delta_net_f32_wmma_head128_inplace": 288, + "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot": 288, + "loom_libs:llm_gated_delta_net_projection_epilogue_f32": 384, + "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32": 384 + }, + "loom_source": "tools/benchmarks/loom/v2_draft_programs.tg_c1.work.loom", + "model": "qwen38_27b_udq4kxl_draft", + "scenario": "tg_c1", + "schema": "ggml-hrx-model-loom-benchmarks-v2" +} 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 new file mode 100644 index 000000000000..3703a6f02ae5 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_draft_programs.tg_c1.work.loom @@ -0,0 +1,2383 @@ +// Generated by tools/benchmarks/generate-model-benchmarks.py for qwen38_27b_udq4kxl_draft scenarios: tg_c1. +// Regenerate from an HRX command program dump rather than editing by hand. + +target.decl @ggml_binary_f32_gfx11_wave64 +target.decl @ggml_binary_swiglu_i4_gfx11_wave32 +target.decl @ggml_copy_f32_gfx11_wave64 +target.decl @ggml_flash_attention_gfx11_wave64 +target.decl @ggml_get_rows_f32_gfx11_wave64 +target.decl @ggml_mul_mat_f32_f32_decode_gfx11_wave64 +target.decl @ggml_mul_mat_gfx11_wave64 +target.decl @ggml_rmsnorm_binary_gfx11_wave32 +target.decl @ggml_rmsnorm_gfx11_wave32 +target.decl @ggml_scale_bias_f32_gfx11_wave64 +target.decl @ggml_set_rows_gfx11_wave64 + +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_binary_f32_gfx11_wave64) @ggml_binary_f32(%element_count: index) launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_binary_swiglu_i4_gfx11_wave32) @ggml_binary_swiglu_symmetric_i4_k32() launch(%lhs: buffer, %rhs: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_concat_dim0_f32() launch(%lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_copy_f32(%element_count: index) launch(%element_count: index, %source: buffer, %output: buffer) +kernel.decl @ggml_fill_negative_f32(%element_count: index) launch(%element_count: index, %output: buffer) +kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) +kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@ggml_mul_mat_f32_f32_decode_gfx11_wave64) @ggml_mul_mat_f32_f32_decode_wave64(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@ggml_mul_mat_gfx11_wave64) @ggml_mul_mat_f32_f32_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_selected_refine_token1(%token_count: index, %candidate_count: index) launch(%token_count: index, %candidate_count: index, %input: buffer, %weight: buffer, %candidates: buffer, %exact_output: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_token1_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_symmetric_i2_scan_token1() launch(%weight: buffer, %output: buffer, %qact: buffer, %scales: buffer, %selected_mask: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma() launch(%first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_quantize_f32_symmetric_i4_k32() launch(%input: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count: index) launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +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_scale_bias_f32_gfx11_wave64) @ggml_scale_bias_f32(%element_count: index) launch(%element_count: index, %input: buffer, %output: buffer) +kernel.decl @ggml_select_symmetric_i4_k32_groups() launch(%input: buffer, %selected_groups: 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 @ggml_top_k128_f32_reduce_gather_register(%element_count: index) launch(%element_count: index, %partial_values: buffer, %partial_ids: buffer, %candidate_output: buffer, %value_output: buffer) +kernel.decl @ggml_top_k8_f32_partitions_register(%element_count: index) launch(%element_count: index, %values: buffer, %partial_values: buffer, %partial_ids: buffer) +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(%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_c1 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_000_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<8xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<8xi8>, tensor<715161600xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_000_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_000_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_001_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_001_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_001_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<40960xi8>, tensor<81920xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>, tensor<81920xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<384xi8> + %second_output = check.generate.fill value(0) : tensor<384xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<384xi8>, tensor<384xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<40960xi8>, tensor<49152xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_004_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_005_ggml_scale_bias_f32_case { + %element_count = check.literal value(30720) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_scale_bias_f32[%element_count](%element_count, %input, %output) : [index](index, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_005_ggml_scale_bias_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_005_ggml_scale_bias_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_006_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(30720) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x30720xf32> + %output = check.generate.fill value(1.0) : tensor<1x30720xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x30720xf32>, tensor<1x30720xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_006_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_006_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_007_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<81920xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<81920xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_007_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_007_llm_ssm_conv_dconv4_silu_rollback_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_008_ggml_scale_bias_f32_case { + %element_count = check.literal value(786432) : index + %input = check.generate.fill value(0) : tensor<3145728xi8> + %output = check.generate.fill value(0) : tensor<3145728xi8> + kernel.launch @ggml_scale_bias_f32[%element_count](%element_count, %input, %output) : [index](index, tensor<3145728xi8>, tensor<3145728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_008_ggml_scale_bias_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_008_ggml_scale_bias_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_009_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(786432) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x786432xf32> + %output = check.generate.fill value(1.0) : tensor<1x786432xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x786432xf32>, tensor<1x786432xf32>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_010_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_010_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_011_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<49152xi8> + %k = check.generate.fill value(0) : tensor<49152xi8> + %v = check.generate.fill value(0) : tensor<65536xi8> + %g = check.generate.fill value(0) : tensor<384xi8> + %beta = check.generate.fill value(0) : tensor<384xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<15728640xi8> + %dst = check.generate.fill value(0) : tensor<49152xi8> + %rms_scales = check.generate.fill value(0) : tensor<384xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<49152xi8>, tensor<49152xi8>, tensor<65536xi8>, tensor<384xi8>, tensor<384xi8>, tensor<3145728xi8>, tensor<15728640xi8>, tensor<49152xi8>, tensor<384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_011_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_draft_tg_c1_011_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(96) : index + %input = check.generate.fill value(0) : tensor<49152xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<49152xi8>, tensor<512xi8>, tensor<49152xi8>, tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_013_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<49152xi8>, tensor<40960xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_013_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_013_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_014_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %residual_output = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<40960xi8>, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_014_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_014_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<139264xi8> + %second_output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_016_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<139264xi8> + %rhs = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<139264xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_016_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_016_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<139264xi8>, tensor<40960xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>, tensor<40960xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<40960xi8>, tensor<98304xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_018_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_019_ggml_mul_mat_f32_f32_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<4300800xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<40960xi8>, tensor<4300800xi8>, tensor<8192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_019_ggml_mul_mat_f32_f32_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_019_ggml_mul_mat_f32_f32_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_020_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<40960xi8>, tensor<8192xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_020_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_020_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_021_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<97280xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<97280xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<49152xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_021_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_021_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_022_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<8192xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<8192xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<8192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_022_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_022_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_023_ggml_set_rows_case { + %token_count = check.literal value(2) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<2x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<2xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<2x1024xf32>, tensor<2xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_023_ggml_set_rows_case> @qwen38_27b_udq4kxl_draft_tg_c1_023_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_024_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(2) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<2x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<2x256xf16> + %output = check.generate.fill value(1.0) : tensor<2x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<2x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<2x256xf16>, tensor<2x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_024_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_024_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_025_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_025_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_025_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_026_ggml_binary_f32_case { + %element_count = check.literal value(10240) : index + %lhs = check.generate.fill value(2.0) : tensor<10240xf32> + %rhs = check.generate.fill value(3.0) : tensor<10240xf32> + %output = check.generate.fill value(0.0) : tensor<10240xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<10240xf32>, tensor<10240xf32>, tensor<10240xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_026_ggml_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_026_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_027_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(2.0) : tensor<2x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<2x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<2x5120xf32>, tensor<5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_027_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_027_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_028_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<2x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_028_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_028_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<20480xi8>, tensor<1042944000xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_030_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<40960xi8>, tensor<40960xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_030_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_030_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_031_ggml_mul_mat_f32_f32_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<43008000xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<43008000xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_031_ggml_mul_mat_f32_f32_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_031_ggml_mul_mat_f32_f32_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_032_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<16xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<16xi8>, tensor<715161600xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_032_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_032_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_033_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_033_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_033_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + %partial = check.generate.fill value(0) : tensor<163840xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>, tensor<163840xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<768xi8> + %second_output = check.generate.fill value(0) : tensor<768xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<768xi8>, tensor<768xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_036_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<81920xi8>, tensor<98304xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_036_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_036_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_037_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<163840xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_038_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_038_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_039_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<131072xi8> + %k = check.generate.fill value(0) : tensor<131072xi8> + %v = check.generate.fill value(0) : tensor<147456xi8> + %g = check.generate.fill value(0) : tensor<768xi8> + %beta = check.generate.fill value(0) : tensor<768xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<40894464xi8> + %dst = check.generate.fill value(0) : tensor<98304xi8> + %rms_scales = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<131072xi8>, tensor<131072xi8>, tensor<147456xi8>, tensor<768xi8>, tensor<768xi8>, tensor<3145728xi8>, tensor<40894464xi8>, tensor<98304xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_039_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_draft_tg_c1_039_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(192) : index + %input = check.generate.fill value(0) : tensor<98304xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<98304xi8>, tensor<512xi8>, tensor<98304xi8>, tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_041_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<98304xi8>, tensor<81920xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_041_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_041_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_042_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %residual_output = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<81920xi8>, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_042_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_042_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<278528xi8> + %second_output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_044_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<278528xi8> + %rhs = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<278528xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_044_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_044_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<278528xi8>, tensor<81920xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>, tensor<81920xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_046_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<196608xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<81920xi8>, tensor<196608xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_046_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_046_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<5611520xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_048_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<81920xi8>, tensor<16384xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_048_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_048_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_049_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<195584xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<195584xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<98304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_049_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_049_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_050_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<16384xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<16384xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_050_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_050_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_051_ggml_set_rows_case { + %token_count = check.literal value(4) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<4x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<4xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<4x1024xf32>, tensor<4xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_051_ggml_set_rows_case> @qwen38_27b_udq4kxl_draft_tg_c1_051_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_052_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<4x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<4x256xf16> + %output = check.generate.fill value(1.0) : tensor<4x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<4x256xf16>, tensor<4x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_052_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_052_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_053_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_053_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_053_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_054_ggml_binary_f32_case { + %element_count = check.literal value(20480) : index + %lhs = check.generate.fill value(2.0) : tensor<20480xf32> + %rhs = check.generate.fill value(3.0) : tensor<20480xf32> + %output = check.generate.fill value(0.0) : tensor<20480xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<20480xf32>, tensor<20480xf32>, tensor<20480xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_054_ggml_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_054_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_055_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(2.0) : tensor<4x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<4x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<4x5120xf32>, tensor<5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_055_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_055_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_056_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<4xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<4x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi32>, tensor<4x5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_056_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_056_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<3973120xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<1042944000xi8>, tensor<3973120xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_058_ggml_get_rows_f32_case { + %token_count = check.literal value(5) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<20xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<20xi8>, tensor<715161600xi8>, tensor<102400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_058_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_058_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_059_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<102400xi8>, tensor<20480xi8>, tensor<102400xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_059_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_059_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + %partial = check.generate.fill value(0) : tensor<204800xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<102400xi8>, tensor<204800xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>, tensor<204800xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<960xi8> + %second_output = check.generate.fill value(0) : tensor<960xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<960xi8>, tensor<960xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_062_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<102400xi8>, tensor<122880xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_062_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_062_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_063_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<204800xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<204800xi8>, tensor<163840xi8>, tensor<204800xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_064_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_064_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_065_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<172032xi8> + %k = check.generate.fill value(0) : tensor<172032xi8> + %v = check.generate.fill value(0) : tensor<188416xi8> + %g = check.generate.fill value(0) : tensor<960xi8> + %beta = check.generate.fill value(0) : tensor<960xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<53477376xi8> + %dst = check.generate.fill value(0) : tensor<122880xi8> + %rms_scales = check.generate.fill value(0) : tensor<960xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<172032xi8>, tensor<172032xi8>, tensor<188416xi8>, tensor<960xi8>, tensor<960xi8>, tensor<3145728xi8>, tensor<53477376xi8>, tensor<122880xi8>, tensor<960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_065_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_draft_tg_c1_065_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(240) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<122880xi8>, tensor<512xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_067_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<122880xi8>, tensor<102400xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_067_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_067_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_068_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(5) : index + %lhs = check.generate.fill value(0) : tensor<102400xi8> + %rhs = check.generate.fill value(0) : tensor<102400xi8> + %residual_output = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<102400xi8>, tensor<102400xi8>, tensor<102400xi8>, tensor<20480xi8>, tensor<102400xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_068_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_068_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<348160xi8> + %second_output = check.generate.fill value(0) : tensor<348160xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<348160xi8>, tensor<348160xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_070_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<348160xi8> + %rhs = check.generate.fill value(0) : tensor<348160xi8> + %output = check.generate.fill value(0) : tensor<348160xi8> + %quantized_values = check.generate.fill value(0) : tensor<43520xi8> + %scales = check.generate.fill value(0) : tensor<10880xi8> + %sums = check.generate.fill value(0) : tensor<10880xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<348160xi8>, tensor<348160xi8>, tensor<348160xi8>, tensor<43520xi8>, tensor<10880xi8>, tensor<10880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_070_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_070_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<348160xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<43520xi8> + %scales = check.generate.fill value(0) : tensor<10880xi8> + %sums = check.generate.fill value(0) : tensor<10880xi8> + %partial = check.generate.fill value(0) : tensor<102400xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<348160xi8>, tensor<102400xi8>, tensor<43520xi8>, tensor<10880xi8>, tensor<10880xi8>, tensor<102400xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_072_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<245760xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<102400xi8>, tensor<245760xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_072_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_072_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<102400xi8>, tensor<5611520xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_074_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<102400xi8>, tensor<20480xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_074_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_074_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_075_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<244736xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<80xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<244736xi8>, tensor<1024xi8>, tensor<80xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_075_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_075_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_076_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<80xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<20480xi8>, tensor<1024xi8>, tensor<80xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_076_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_076_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_077_ggml_set_rows_case { + %token_count = check.literal value(5) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<5x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<5xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<5x1024xf32>, tensor<5xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_077_ggml_set_rows_case> @qwen38_27b_udq4kxl_draft_tg_c1_077_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_078_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(5) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<5x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<5x256xf16> + %output = check.generate.fill value(1.0) : tensor<5x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<5x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<5x256xf16>, tensor<5x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_078_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_078_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_079_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<122880xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_079_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_079_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_080_ggml_binary_f32_case { + %element_count = check.literal value(25600) : index + %lhs = check.generate.fill value(2.0) : tensor<25600xf32> + %rhs = check.generate.fill value(3.0) : tensor<25600xf32> + %output = check.generate.fill value(0.0) : tensor<25600xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<25600xf32>, tensor<25600xf32>, tensor<25600xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_080_ggml_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_080_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_081_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(2.0) : tensor<5x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<5x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<5x5120xf32>, tensor<5120xf32>, tensor<5x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_081_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_081_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_082_ggml_get_rows_f32_case { + %token_count = check.literal value(5) : index + %row_count = check.literal value(5) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<5xi32> + %weight = check.generate.fill value(0.0) : tensor<5x5120xf32> + %output = check.generate.fill value(1.0) : tensor<5x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<5xi32>, tensor<5x5120xf32>, tensor<5x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_082_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_082_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<4966400xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<102400xi8>, tensor<1042944000xi8>, tensor<4966400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_084_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<102400xi8> + %rhs = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<102400xi8>, tensor<102400xi8>, tensor<204800xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_084_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_084_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<204800xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<204800xi8>, tensor<56115200xi8>, tensor<102400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_086_ggml_get_rows_f32_case { + %token_count = check.literal value(3) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<12xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<12xi8>, tensor<715161600xi8>, tensor<61440xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_086_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_086_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_087_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<61440xi8>, tensor<20480xi8>, tensor<61440xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_087_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_087_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + %partial = check.generate.fill value(0) : tensor<122880xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<61440xi8>, tensor<122880xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>, tensor<122880xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<576xi8> + %second_output = check.generate.fill value(0) : tensor<576xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<576xi8>, tensor<576xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_090_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<73728xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<61440xi8>, tensor<73728xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_090_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_090_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_091_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<122880xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<122880xi8>, tensor<163840xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_092_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_092_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_093_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<90112xi8> + %k = check.generate.fill value(0) : tensor<90112xi8> + %v = check.generate.fill value(0) : tensor<106496xi8> + %g = check.generate.fill value(0) : tensor<576xi8> + %beta = check.generate.fill value(0) : tensor<576xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<28311552xi8> + %dst = check.generate.fill value(0) : tensor<73728xi8> + %rms_scales = check.generate.fill value(0) : tensor<576xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<90112xi8>, tensor<90112xi8>, tensor<106496xi8>, tensor<576xi8>, tensor<576xi8>, tensor<3145728xi8>, tensor<28311552xi8>, tensor<73728xi8>, tensor<576xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_093_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_draft_tg_c1_093_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(144) : index + %input = check.generate.fill value(0) : tensor<73728xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<73728xi8> + %output = check.generate.fill value(0) : tensor<73728xi8> + %quantized_values = check.generate.fill value(0) : tensor<9216xi8> + %scales = check.generate.fill value(0) : tensor<2304xi8> + %sums = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<73728xi8>, tensor<512xi8>, tensor<73728xi8>, tensor<73728xi8>, tensor<9216xi8>, tensor<2304xi8>, tensor<2304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_095_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<73728xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<9216xi8> + %scales = check.generate.fill value(0) : tensor<2304xi8> + %sums = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<73728xi8>, tensor<61440xi8>, tensor<9216xi8>, tensor<2304xi8>, tensor<2304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_095_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_095_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_096_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(3) : index + %lhs = check.generate.fill value(0) : tensor<61440xi8> + %rhs = check.generate.fill value(0) : tensor<61440xi8> + %residual_output = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<61440xi8>, tensor<61440xi8>, tensor<61440xi8>, tensor<20480xi8>, tensor<61440xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_096_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_096_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<208896xi8> + %second_output = check.generate.fill value(0) : tensor<208896xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<208896xi8>, tensor<208896xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_098_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<208896xi8> + %rhs = check.generate.fill value(0) : tensor<208896xi8> + %output = check.generate.fill value(0) : tensor<208896xi8> + %quantized_values = check.generate.fill value(0) : tensor<26112xi8> + %scales = check.generate.fill value(0) : tensor<6528xi8> + %sums = check.generate.fill value(0) : tensor<6528xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<208896xi8>, tensor<208896xi8>, tensor<208896xi8>, tensor<26112xi8>, tensor<6528xi8>, tensor<6528xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_098_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_098_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<208896xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<26112xi8> + %scales = check.generate.fill value(0) : tensor<6528xi8> + %sums = check.generate.fill value(0) : tensor<6528xi8> + %partial = check.generate.fill value(0) : tensor<61440xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<208896xi8>, tensor<61440xi8>, tensor<26112xi8>, tensor<6528xi8>, tensor<6528xi8>, tensor<61440xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_100_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<147456xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<61440xi8>, tensor<147456xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_100_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_100_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<61440xi8>, tensor<5611520xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_102_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<12288xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<61440xi8>, tensor<12288xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_102_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_102_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_103_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<146432xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<48xi8> + %output = check.generate.fill value(0) : tensor<73728xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<146432xi8>, tensor<1024xi8>, tensor<48xi8>, tensor<73728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_103_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_103_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_104_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<12288xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<48xi8> + %output = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<12288xi8>, tensor<1024xi8>, tensor<48xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_104_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_104_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_105_ggml_set_rows_case { + %token_count = check.literal value(3) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<3x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<3xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<3x1024xf32>, tensor<3xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_105_ggml_set_rows_case> @qwen38_27b_udq4kxl_draft_tg_c1_105_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_106_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(3) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<3x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<3x256xf16> + %output = check.generate.fill value(1.0) : tensor<3x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<3x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<3x256xf16>, tensor<3x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_106_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_106_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_107_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<73728xi8> + %quantized_values = check.generate.fill value(0) : tensor<9216xi8> + %scales = check.generate.fill value(0) : tensor<2304xi8> + %sums = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<73728xi8>, tensor<9216xi8>, tensor<2304xi8>, tensor<2304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_107_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_107_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_108_ggml_binary_f32_case { + %element_count = check.literal value(15360) : index + %lhs = check.generate.fill value(2.0) : tensor<15360xf32> + %rhs = check.generate.fill value(3.0) : tensor<15360xf32> + %output = check.generate.fill value(0.0) : tensor<15360xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<15360xf32>, tensor<15360xf32>, tensor<15360xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_108_ggml_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_108_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_109_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(2.0) : tensor<3x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<3x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<3x5120xf32>, tensor<5120xf32>, tensor<3x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_109_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_109_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_110_ggml_get_rows_f32_case { + %token_count = check.literal value(3) : index + %row_count = check.literal value(3) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<3xi32> + %weight = check.generate.fill value(0.0) : tensor<3x5120xf32> + %output = check.generate.fill value(1.0) : tensor<3x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<3xi32>, tensor<3x5120xf32>, tensor<3x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_110_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_110_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<2979840xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<61440xi8>, tensor<1042944000xi8>, tensor<2979840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_112_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<61440xi8> + %rhs = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<61440xi8>, tensor<61440xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_112_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_112_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<122880xi8>, tensor<56115200xi8>, tensor<61440xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_114_ggml_get_rows_f32_case { + %token_count = check.literal value(16) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<64xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<64xi8>, tensor<715161600xi8>, tensor<327680xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_114_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_114_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_115_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<327680xi8>, tensor<20480xi8>, tensor<327680xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_115_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_115_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + %partial = check.generate.fill value(0) : tensor<655360xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<327680xi8>, tensor<655360xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>, tensor<655360xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<3072xi8> + %second_output = check.generate.fill value(0) : tensor<3072xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<3072xi8>, tensor<3072xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_118_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<327680xi8>, tensor<393216xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_118_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_118_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_119_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<655360xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<655360xi8>, tensor<163840xi8>, tensor<655360xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_120_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_120_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_121_llm_gated_delta_net_f32_wmma_head128_case { + %q = check.generate.fill value(0) : tensor<417792xi8> + %k = check.generate.fill value(0) : tensor<417792xi8> + %v = check.generate.fill value(0) : tensor<434176xi8> + %g = check.generate.fill value(0) : tensor<2112xi8> + %beta = check.generate.fill value(0) : tensor<2112xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<3416064xi8> + %rms_scales = check.generate.fill value(0) : tensor<2112xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128[](%q, %k, %v, %g, %beta, %state_in, %dst, %rms_scales) : [](tensor<417792xi8>, tensor<417792xi8>, tensor<434176xi8>, tensor<2112xi8>, tensor<2112xi8>, tensor<3145728xi8>, tensor<3416064xi8>, tensor<2112xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_121_llm_gated_delta_net_f32_wmma_head128_case> @qwen38_27b_udq4kxl_draft_tg_c1_121_llm_gated_delta_net_f32_wmma_head128 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_122_ggml_copy_f32_case { + %element_count = check.literal value(786432) : index + %source = check.generate.fill value(0) : tensor<3145728xi8> + %output = check.generate.fill value(0) : tensor<3145728xi8> + kernel.launch @ggml_copy_f32[%element_count](%element_count, %source, %output) : [index](index, tensor<3145728xi8>, tensor<3145728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_122_ggml_copy_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_122_ggml_copy_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_123_llm_gated_delta_net_f32_wmma_head128_inplace_case { + %q = check.generate.fill value(0) : tensor<8192xi8> + %k = check.generate.fill value(0) : tensor<8192xi8> + %v = check.generate.fill value(0) : tensor<24576xi8> + %g = check.generate.fill value(0) : tensor<192xi8> + %beta = check.generate.fill value(0) : tensor<192xi8> + %state_inout = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<24576xi8> + %rms_scales = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_inplace[](%q, %k, %v, %g, %beta, %state_inout, %dst, %rms_scales) : [](tensor<8192xi8>, tensor<8192xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3145728xi8>, tensor<24576xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_123_llm_gated_delta_net_f32_wmma_head128_inplace_case> @qwen38_27b_udq4kxl_draft_tg_c1_123_llm_gated_delta_net_f32_wmma_head128_inplace + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(768) : index + %input = check.generate.fill value(0) : tensor<393216xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<393216xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<393216xi8>, tensor<512xi8>, tensor<393216xi8>, tensor<393216xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_125_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<393216xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<393216xi8>, tensor<327680xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_125_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_125_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_126_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(16) : index + %lhs = check.generate.fill value(0) : tensor<327680xi8> + %rhs = check.generate.fill value(0) : tensor<327680xi8> + %residual_output = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<327680xi8>, tensor<327680xi8>, tensor<327680xi8>, tensor<20480xi8>, tensor<327680xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_126_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_126_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<1114112xi8> + %second_output = check.generate.fill value(0) : tensor<1114112xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<1114112xi8>, tensor<1114112xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_128_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<1114112xi8> + %rhs = check.generate.fill value(0) : tensor<1114112xi8> + %output = check.generate.fill value(0) : tensor<1114112xi8> + %quantized_values = check.generate.fill value(0) : tensor<139264xi8> + %scales = check.generate.fill value(0) : tensor<34816xi8> + %sums = check.generate.fill value(0) : tensor<34816xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<1114112xi8>, tensor<1114112xi8>, tensor<1114112xi8>, tensor<139264xi8>, tensor<34816xi8>, tensor<34816xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_128_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_128_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<1114112xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<139264xi8> + %scales = check.generate.fill value(0) : tensor<34816xi8> + %sums = check.generate.fill value(0) : tensor<34816xi8> + %partial = check.generate.fill value(0) : tensor<327680xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<1114112xi8>, tensor<327680xi8>, tensor<139264xi8>, tensor<34816xi8>, tensor<34816xi8>, tensor<327680xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_130_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<786432xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<327680xi8>, tensor<786432xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_130_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_130_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<327680xi8>, tensor<5611520xi8>, tensor<65536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_132_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<327680xi8>, tensor<65536xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_132_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_132_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_133_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<785408xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<785408xi8>, tensor<1024xi8>, tensor<256xi8>, tensor<393216xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_133_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_133_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_134_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<65536xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<65536xi8>, tensor<1024xi8>, tensor<256xi8>, tensor<65536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_134_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_134_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_135_ggml_set_rows_case { + %token_count = check.literal value(16) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<16x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<16xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<16x1024xf32>, tensor<16xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_135_ggml_set_rows_case> @qwen38_27b_udq4kxl_draft_tg_c1_135_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_136_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(16) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<16x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<16x256xf16> + %output = check.generate.fill value(1.0) : tensor<16x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<16x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<16x256xf16>, tensor<16x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_136_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_136_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_137_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<393216xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_137_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_137_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_138_ggml_binary_f32_case { + %element_count = check.literal value(81920) : index + %lhs = check.generate.fill value(2.0) : tensor<81920xf32> + %rhs = check.generate.fill value(3.0) : tensor<81920xf32> + %output = check.generate.fill value(0.0) : tensor<81920xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<81920xf32>, tensor<81920xf32>, tensor<81920xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_138_ggml_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_138_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_139_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(2.0) : tensor<16x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<16x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<16x5120xf32>, tensor<5120xf32>, tensor<16x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_139_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_139_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_140_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<327680xi8> + %rhs = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<327680xi8>, tensor<327680xi8>, tensor<655360xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_140_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_140_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<655360xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<655360xi8>, tensor<56115200xi8>, tensor<327680xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_142_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<4x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_142_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_142_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_143_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<81920xi8>, tensor<81920xi8>, tensor<163840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_143_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_143_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<163840xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<163840xi8>, tensor<56115200xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_145_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<4xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi8>, tensor<715161600xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_145_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_145_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_146_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<1x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<1x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<1x5120xf32>, tensor<5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_146_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_146_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_147_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<20480xi8> + %rhs = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<20480xi8>, tensor<20480xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_147_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_147_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_148_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(10240) : index + %output_size = check.literal value(5120) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<43008000xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<40960xi8>, tensor<43008000xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_148_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_draft_tg_c1_148_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_149_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_149_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_149_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_150_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<20480xi8>, tensor<49152xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_150_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_150_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_151_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(5120) : index + %output_size = check.literal value(1024) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<4300800xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<20480xi8>, tensor<4300800xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_151_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_draft_tg_c1_151_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_152_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<20480xi8>, tensor<4096xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_152_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_152_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_153_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<48128xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<16xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<48128xi8>, tensor<1024xi8>, tensor<16xi8>, tensor<24576xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_153_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_153_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_154_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<4096xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<16xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<4096xi8>, tensor<1024xi8>, tensor<16xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_154_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_154_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_155_ggml_set_rows_case { + %token_count = check.literal value(1) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<1x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<1xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<1x1024xf32>, tensor<1xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_155_ggml_set_rows_case> @qwen38_27b_udq4kxl_draft_tg_c1_155_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_156_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<1x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<1x256xf16> + %output = check.generate.fill value(1.0) : tensor<1x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<1x256xf16>, tensor<1x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_156_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_156_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_157_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<24576xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_157_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_157_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_158_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<24576xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<24576xi8>, tensor<20480xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_158_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_158_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_159_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(1) : index + %lhs = check.generate.fill value(0) : tensor<20480xi8> + %rhs = check.generate.fill value(0) : tensor<20480xi8> + %residual_output = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_159_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_159_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<69632xi8> + %second_output = check.generate.fill value(0) : tensor<69632xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<69632xi8>, tensor<69632xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_161_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<69632xi8> + %rhs = check.generate.fill value(0) : tensor<69632xi8> + %output = check.generate.fill value(0) : tensor<69632xi8> + %quantized_values = check.generate.fill value(0) : tensor<8704xi8> + %scales = check.generate.fill value(0) : tensor<2176xi8> + %sums = check.generate.fill value(0) : tensor<2176xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<69632xi8>, tensor<69632xi8>, tensor<69632xi8>, tensor<8704xi8>, tensor<2176xi8>, tensor<2176xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_161_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_161_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<69632xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<8704xi8> + %scales = check.generate.fill value(0) : tensor<2176xi8> + %sums = check.generate.fill value(0) : tensor<2176xi8> + %partial = check.generate.fill value(0) : tensor<20480xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<69632xi8>, tensor<20480xi8>, tensor<8704xi8>, tensor<2176xi8>, tensor<2176xi8>, tensor<20480xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_163_ggml_binary_f32_case { + %element_count = check.literal value(5120) : index + %lhs = check.generate.fill value(2.0) : tensor<5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<5120xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<5120xf32>, tensor<5120xf32>, tensor<5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_163_ggml_binary_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_163_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_164_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(1) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<1x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<1x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_164_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_164_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_165_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_165_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_165_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_166_ggml_mul_mat_q6_k_symmetric_i2_scan_token1_case { + %weight = check.generate.fill value(0) : tensor<1380659200xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + %qact = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %selected_mask = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_q6_k_symmetric_i2_scan_token1[](%weight, %output, %qact, %scales, %selected_mask) : [](tensor<1380659200xi8>, tensor<993280xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_166_ggml_mul_mat_q6_k_symmetric_i2_scan_token1_case> @qwen38_27b_udq4kxl_draft_tg_c1_166_ggml_mul_mat_q6_k_symmetric_i2_scan_token1 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_167_ggml_top_k8_f32_partitions_register_case { + %element_count = check.literal value(248320) : index + %values = check.generate.fill value(0) : tensor<993280xi8> + %partial_values = check.generate.fill value(0) : tensor<4096xi8> + %partial_ids = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_top_k8_f32_partitions_register[%element_count](%element_count, %values, %partial_values, %partial_ids) : [index](index, tensor<993280xi8>, tensor<4096xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_167_ggml_top_k8_f32_partitions_register_case> @qwen38_27b_udq4kxl_draft_tg_c1_167_ggml_top_k8_f32_partitions_register + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_168_ggml_top_k128_f32_reduce_gather_register_case { + %element_count = check.literal value(248320) : index + %partial_values = check.generate.fill value(0) : tensor<4096xi8> + %partial_ids = check.generate.fill value(0) : tensor<4096xi8> + %candidate_output = check.generate.fill value(0) : tensor<512xi8> + %value_output = check.generate.fill value(0) : tensor<512xi8> + kernel.launch @ggml_top_k128_f32_reduce_gather_register[%element_count](%element_count, %partial_values, %partial_ids, %candidate_output, %value_output) : [index](index, tensor<4096xi8>, tensor<4096xi8>, tensor<512xi8>, tensor<512xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_168_ggml_top_k128_f32_reduce_gather_register_case> @qwen38_27b_udq4kxl_draft_tg_c1_168_ggml_top_k128_f32_reduce_gather_register + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_169_ggml_fill_negative_f32_case { + %element_count = check.literal value(248320) : index + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_fill_negative_f32[%element_count](%element_count, %output) : [index](index, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_169_ggml_fill_negative_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_169_ggml_fill_negative_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_170_ggml_mul_mat_q6_k_packed_selected_refine_token1_case { + %token_count = check.literal value(1) : index + %candidate_count = check.literal value(64) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1380659200xi8> + %candidates = check.generate.fill value(0) : tensor<256xi8> + %exact_output = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_selected_refine_token1[%token_count, %candidate_count](%token_count, %candidate_count, %input, %weight, %candidates, %exact_output, %output) : [index, index](index, index, tensor<20480xi8>, tensor<1380659200xi8>, tensor<256xi8>, tensor<256xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_170_ggml_mul_mat_q6_k_packed_selected_refine_token1_case> @qwen38_27b_udq4kxl_draft_tg_c1_170_ggml_mul_mat_q6_k_packed_selected_refine_token1 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_171_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>, tensor<40960xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_171_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_171_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_172_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(5120) : index + %output_size = check.literal value(48) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<138240xi8> + %output = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<20480xi8>, tensor<138240xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_172_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_draft_tg_c1_172_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_173_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<20480xi8>, tensor<24576xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_173_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_173_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_174_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<40960xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<40960xi8>, tensor<163840xi8>, tensor<40960xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_175_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_175_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_176_llm_gated_delta_net_f32_wmma_head128_inplace_case { + %q = check.generate.fill value(0) : tensor<8192xi8> + %k = check.generate.fill value(0) : tensor<8192xi8> + %v = check.generate.fill value(0) : tensor<24576xi8> + %g = check.generate.fill value(0) : tensor<192xi8> + %beta = check.generate.fill value(0) : tensor<192xi8> + %state_inout = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<24576xi8> + %rms_scales = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_inplace[](%q, %k, %v, %g, %beta, %state_inout, %dst, %rms_scales) : [](tensor<8192xi8>, tensor<8192xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3145728xi8>, tensor<24576xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_176_llm_gated_delta_net_f32_wmma_head128_inplace_case> @qwen38_27b_udq4kxl_draft_tg_c1_176_llm_gated_delta_net_f32_wmma_head128_inplace + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_177_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(48) : index + %input = check.generate.fill value(0) : tensor<24576xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<24576xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<24576xi8>, tensor<512xi8>, tensor<24576xi8>, tensor<24576xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_177_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_draft_tg_c1_177_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_178_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<2xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<2x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<2xi32>, tensor<2x5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_178_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_draft_tg_c1_178_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_draft_tg_c1_179_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<1986560xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<40960xi8>, tensor<1042944000xi8>, tensor<1986560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_draft_tg_c1_179_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_draft_tg_c1_179_ggml_mul_mat_q6_k_packed_token1_f16_wmma 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 new file mode 100644 index 000000000000..c25c23af3c4d --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.tg_c5.json @@ -0,0 +1,5812 @@ +{ + "command_count": 8673, + "dispatch_count": 181, + "dispatches": [ + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_000_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "2", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 2 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_005_ggml_scale_bias_f32", + "compile_parameters": { + "ggml.scale.bias": "0", + "ggml.scale.scale": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "element_count": 30720 + }, + "kernel": "loom_libs:ggml_scale_bias_f32", + "library_sources": [], + "primary_sources": [ + "ops/scale_bias_f32.loom" + ], + "sources": [ + "ops/scale_bias_f32.loom" + ], + "symbol": "ggml_scale_bias_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_006_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "30720", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 384, + "integer_parameters": { + "hidden_size": 30720, + "row_count": 20, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "2", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_008_ggml_scale_bias_f32", + "compile_parameters": { + "ggml.scale.bias": "0", + "ggml.scale.scale": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "element_count": 786432 + }, + "kernel": "loom_libs:ggml_scale_bias_f32", + "library_sources": [], + "primary_sources": [ + "ops/scale_bias_f32.loom" + ], + "sources": [ + "ops/scale_bias_f32.loom" + ], + "symbol": "ggml_scale_bias_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_009_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "786432", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 384, + "integer_parameters": { + "hidden_size": 786432, + "row_count": 20, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 96 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "20480", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "96", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "2", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "20480", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "token_count": 96 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 256, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 130, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_019_ggml_mul_mat_f32_f32_wmma", + "compile_parameters": { + "ggml.mul_mat.input_size": "5120", + "ggml.mul_mat.output_accumulation": "0", + "ggml.mul_mat.output_size": "1024", + "ggml.mul_mat.output_unary_op": "23", + "ggml.mul_mat.weight_format": "6", + "ggml.workload.token_capacity": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_wmma.loom", + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_021_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "24576", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "2", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "12288", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_022_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "2048", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "2", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "2048", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_023_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 68, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 2 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_024_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 2 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_026_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "element_count": 10240 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_027_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 8, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_028_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 2, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_030_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_031_ggml_mul_mat_f32_f32_wmma", + "compile_parameters": { + "ggml.mul_mat.input_size": "10240", + "ggml.mul_mat.output_accumulation": "0", + "ggml.mul_mat.output_size": "5120", + "ggml.mul_mat.output_unary_op": "23", + "ggml.mul_mat.weight_format": "6", + "ggml.workload.token_capacity": "2" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_wmma", + "library_sources": [ + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_wmma.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_wmma.loom", + "motifs/mul_mat_f32_f32_wmma_core.loom", + "motifs/dequant.loom", + "motifs/unary_f32_apply.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_032_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "4", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "4", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 192 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "40960", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "192", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "4", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "40960", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "token_count": 192 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 255, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_049_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "49152", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "4", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "24576", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_050_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "4096", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "4", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "4096", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_051_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_052_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 4 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 33, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_054_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "element_count": 20480 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_055_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 5, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_056_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "4", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 4, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_058_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "5", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 5 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_059_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_062_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_063_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "5", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_064_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 240 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_065_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "51200", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "240", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "5", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "51200", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 240 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_067_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_068_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 128, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_070_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_072_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_074_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_075_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "61440", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "5", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "30720", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_076_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "5120", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "5", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "5120", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_077_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 5 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_078_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 5 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_079_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_080_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "element_count": 25600 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_081_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_082_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "5", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 5, + "token_count": 5 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_084_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "5" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 5 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_086_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "3", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 3 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_087_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_090_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_091_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "3", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_092_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 144 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_093_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "30720", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "144", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "3", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "30720", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 144 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_095_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_096_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 128, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_098_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_100_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_102_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_103_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "36864", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "3", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "18432", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_104_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "3072", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "3", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "3072", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_105_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 3 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_106_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 3 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_107_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_108_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "element_count": 15360 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_109_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_110_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "3", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 3, + "token_count": 3 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_112_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "3" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 3 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_114_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "16", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 16 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_115_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_118_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_119_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "16", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_120_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 768 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_121_llm_gated_delta_net_f32_wmma_head128", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "163840", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "768", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.token_count": "11", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "163840", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_122_ggml_copy_f32", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 288, + "integer_parameters": { + "element_count": 786432 + }, + "kernel": "loom_libs:ggml_copy_f32", + "library_sources": [ + "motifs/copy/vector_load_store.loom" + ], + "primary_sources": [ + "ops/copy/copy_f32.loom" + ], + "sources": [ + "ops/copy/copy_f32.loom", + "motifs/copy/vector_load_store.loom" + ], + "symbol": "ggml_copy_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_123_llm_gated_delta_net_f32_wmma_head128_inplace", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "163840", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "768", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.token_count": "1", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "163840", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 240, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_inplace", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_inplace", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 768 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_125_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_126_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 128, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_128_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 65, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_133_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "196608", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "16", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "98304", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_134_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "16384", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "16", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "16384", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_135_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 34, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 16 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_136_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 16 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 17, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_138_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "element_count": 81920 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_139_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 4, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_140_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "16" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 16 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_142_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 4, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_143_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "10240", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "5120" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_145_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_146_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 7, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_147_ggml_concat_dim0_f32", + "compile_parameters": { + "ggml.concat_dim0_f32.lhs_width": "5120", + "ggml.concat_dim0_f32.rhs_width": "5120", + "ggml.concat_dim0_f32.row_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_concat_dim0_f32", + "library_sources": [ + "motifs/copy/concat_dim0.loom" + ], + "primary_sources": [ + "ops/copy/concat_dim0_f32.loom" + ], + "sources": [ + "ops/copy/concat_dim0_f32.loom", + "motifs/copy/concat_dim0.loom" + ], + "symbol": "ggml_concat_dim0_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_148_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "5120", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "6" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "input_size": 10240, + "output_size": 5120, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_149_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_150_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_151_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "1024", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "6" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": { + "input_size": 5120, + "output_size": 1024, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_152_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_153_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "12288", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "1", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "6144", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_154_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "1024", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "1", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "1024", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_155_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 36, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_156_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 1 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_157_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 18, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_158_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_159_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 129, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_161_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 66, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_163_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 3, + "integer_parameters": { + "element_count": 5120 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_164_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 1, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_165_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "5120", + "ggml.quantize_symmetric_i4_k32.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_166_ggml_select_symmetric_i4_k32_groups", + "compile_parameters": { + "ggml.mul_mat_q6_k_shortlist.input_size": "5120", + "ggml.mul_mat_q6_k_shortlist.selected_group_count": "96" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_select_symmetric_i4_k32_groups", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_select_symmetric_i4_k32_groups", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_167_ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "compile_parameters": { + "ggml.mul_mat_q6_k_shortlist.input_size": "5120", + "ggml.mul_mat_q6_k_shortlist.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_symmetric_i2_scan_token1", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_168_ggml_top_k8_f32_partitions_register", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 248320 + }, + "kernel": "loom_libs:ggml_top_k8_f32_partitions_register", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_top_k8_f32_partitions_register", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_169_ggml_top_k128_f32_reduce_gather_register", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 248320 + }, + "kernel": "loom_libs:ggml_top_k128_f32_reduce_gather_register", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_top_k128_f32_reduce_gather_register", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_170_ggml_fill_negative_f32", + "compile_parameters": {}, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 248320 + }, + "kernel": "loom_libs:ggml_fill_negative_f32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_fill_negative_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_171_ggml_mul_mat_q6_k_packed_selected_refine_token1", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320", + "ggml.mul_mat_q6_k_packed.weight_offset": "337715200" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 2, + "integer_parameters": { + "candidate_count": 64, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_selected_refine_token1", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_selected_refine_token1", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "candidate_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_172_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_173_ggml_mul_mat_f32_f32_decode_wave64", + "compile_parameters": { + "ggml.mul_mat_f32_f32_decode.output_capacity": "48", + "ggml.mul_mat_f32_f32_decode.token_capacity": "1", + "ggml.mul_mat_f32_f32_decode.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 96, + "integer_parameters": { + "input_size": 5120, + "output_size": 48, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + "library_sources": [ + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/mul_mat_f32_f32_decode.loom" + ], + "sources": [ + "ops/mul_mat_f32_f32_decode.loom", + "motifs/dequant.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_mul_mat_f32_f32_decode_wave64", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "input_size", + "type": "index" + }, + { + "name": "output_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_174_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "1" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_175_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "1", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_176_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 48 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_177_llm_gated_delta_net_f32_wmma_head128_inplace", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "10240", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "48", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.token_count": "1", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "10240", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_inplace", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_inplace", + "workload_parameters": [] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_178_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 48 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_179_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "2", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 2, + "token_count": 2 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@qwen38_27b_udq4kxl_companion_tg_c5_180_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 2 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + } + ], + "generated_count": 181, + "kernel_counts": { + "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32": 1024, + "loom_libs:ggml_binary_f32": 16, + "loom_libs:ggml_binary_swiglu_symmetric_i4_k32": 520, + "loom_libs:ggml_concat_dim0_f32": 8, + "loom_libs:ggml_copy_f32": 288, + "loom_libs:ggml_fill_negative_f32": 1, + "loom_libs:ggml_flash_attention_f32_f16_wmma": 136, + "loom_libs:ggml_get_rows_f32": 793, + "loom_libs:ggml_mul_mat_f32_f32_decode_wave64": 116, + "loom_libs:ggml_mul_mat_f32_f32_wmma": 36, + "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma": 88, + "loom_libs:ggml_mul_mat_q6_k_packed_selected_refine_token1": 2, + "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma": 8, + "loom_libs:ggml_mul_mat_q6_k_symmetric_i2_scan_token1": 1, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma": 856, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma": 904, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma": 1176, + "loom_libs:ggml_quantize_f32_symmetric_i4_k32": 137, + "loom_libs:ggml_rmsnorm_binary_f32": 32, + "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32": 16, + "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32": 384, + "loom_libs:ggml_rmsnorm_mul_rope_f32": 272, + "loom_libs:ggml_scale_bias_f32": 192, + "loom_libs:ggml_select_symmetric_i4_k32_groups": 1, + "loom_libs:ggml_set_rows": 272, + "loom_libs:ggml_top_k128_f32_reduce_gather_register": 1, + "loom_libs:ggml_top_k8_f32_partitions_register": 1, + "loom_libs:llm_gated_delta_net_f32_wmma_head128": 48, + "loom_libs:llm_gated_delta_net_f32_wmma_head128_inplace": 288, + "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot": 288, + "loom_libs:llm_gated_delta_net_projection_epilogue_f32": 384, + "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32": 384 + }, + "loom_source": "tools/benchmarks/loom/v2_mtp4_companion_all.work.loom", + "model": "qwen38_27b_udq4kxl_companion", + "scenario": "tg_c5", + "schema": "ggml-hrx-model-loom-benchmarks-v2" +} 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 new file mode 100644 index 000000000000..e92702b54836 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_companion_all.work.loom @@ -0,0 +1,2392 @@ +// Generated by tools/benchmarks/generate-model-benchmarks.py for qwen38_27b_udq4kxl_companion scenarios: tg_c5. +// Regenerate from an HRX command program dump rather than editing by hand. + +target.decl @ggml_binary_f32_gfx11_wave64 +target.decl @ggml_binary_swiglu_i4_gfx11_wave32 +target.decl @ggml_copy_f32_gfx11_wave64 +target.decl @ggml_flash_attention_gfx11_wave64 +target.decl @ggml_get_rows_f32_gfx11_wave64 +target.decl @ggml_mul_mat_f32_f32_decode_gfx11_wave64 +target.decl @ggml_mul_mat_gfx11_wave64 +target.decl @ggml_rmsnorm_binary_gfx11_wave32 +target.decl @ggml_rmsnorm_gfx11_wave32 +target.decl @ggml_scale_bias_f32_gfx11_wave64 +target.decl @ggml_set_rows_gfx11_wave64 + +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_binary_f32_gfx11_wave64) @ggml_binary_f32(%element_count: index) launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_binary_swiglu_i4_gfx11_wave32) @ggml_binary_swiglu_symmetric_i4_k32() launch(%lhs: buffer, %rhs: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_concat_dim0_f32() launch(%lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_copy_f32_gfx11_wave64) @ggml_copy_f32(%element_count: index) launch(%element_count: index, %source: buffer, %output: buffer) +kernel.decl @ggml_fill_negative_f32(%element_count: index) launch(%element_count: index, %output: buffer) +kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) +kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@ggml_mul_mat_f32_f32_decode_gfx11_wave64) @ggml_mul_mat_f32_f32_decode_wave64(%token_count: index, %input_size: index, %output_size: index) launch(%token_count: index, %input_size: index, %output_size: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl target(@ggml_mul_mat_gfx11_wave64) @ggml_mul_mat_f32_f32_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_selected_refine_token1(%token_count: index, %candidate_count: index) launch(%token_count: index, %candidate_count: index, %input: buffer, %weight: buffer, %candidates: buffer, %exact_output: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_token1_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_symmetric_i2_scan_token1() launch(%weight: buffer, %output: buffer, %qact: buffer, %scales: buffer, %selected_mask: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma() launch(%first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_quantize_f32_symmetric_i4_k32() launch(%input: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count: index) launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +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_scale_bias_f32_gfx11_wave64) @ggml_scale_bias_f32(%element_count: index) launch(%element_count: index, %input: buffer, %output: buffer) +kernel.decl @ggml_select_symmetric_i4_k32_groups() launch(%input: buffer, %selected_mask: 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 @ggml_top_k128_f32_reduce_gather_register(%element_count: index) launch(%element_count: index, %partial_values: buffer, %partial_ids: buffer, %candidate_output: buffer, %value_output: buffer) +kernel.decl @ggml_top_k8_f32_partitions_register(%element_count: index) launch(%element_count: index, %values: buffer, %partial_values: buffer, %partial_ids: buffer) +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(%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 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_000_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<8xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<8xi8>, tensor<715161600xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_000_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_000_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_001_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<40960xi8>, tensor<81920xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>, tensor<81920xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<384xi8> + %second_output = check.generate.fill value(0) : tensor<384xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<384xi8>, tensor<384xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<40960xi8>, tensor<49152xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_004_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_005_ggml_scale_bias_f32_case { + %element_count = check.literal value(30720) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_scale_bias_f32[%element_count](%element_count, %input, %output) : [index](index, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_005_ggml_scale_bias_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_005_ggml_scale_bias_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_006_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(30720) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x30720xf32> + %output = check.generate.fill value(1.0) : tensor<1x30720xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x30720xf32>, tensor<1x30720xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_006_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_006_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<81920xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<81920xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_007_llm_ssm_conv_dconv4_silu_rollback_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_008_ggml_scale_bias_f32_case { + %element_count = check.literal value(786432) : index + %input = check.generate.fill value(0) : tensor<3145728xi8> + %output = check.generate.fill value(0) : tensor<3145728xi8> + kernel.launch @ggml_scale_bias_f32[%element_count](%element_count, %input, %output) : [index](index, tensor<3145728xi8>, tensor<3145728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_008_ggml_scale_bias_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_008_ggml_scale_bias_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_009_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(786432) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x786432xf32> + %output = check.generate.fill value(1.0) : tensor<1x786432xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x786432xf32>, tensor<1x786432xf32>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_010_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<49152xi8> + %k = check.generate.fill value(0) : tensor<49152xi8> + %v = check.generate.fill value(0) : tensor<65536xi8> + %g = check.generate.fill value(0) : tensor<384xi8> + %beta = check.generate.fill value(0) : tensor<384xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<15728640xi8> + %dst = check.generate.fill value(0) : tensor<49152xi8> + %rms_scales = check.generate.fill value(0) : tensor<384xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<49152xi8>, tensor<49152xi8>, tensor<65536xi8>, tensor<384xi8>, tensor<384xi8>, tensor<3145728xi8>, tensor<15728640xi8>, tensor<49152xi8>, tensor<384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_companion_tg_c5_011_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(96) : index + %input = check.generate.fill value(0) : tensor<49152xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<49152xi8>, tensor<512xi8>, tensor<49152xi8>, tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_012_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<49152xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<49152xi8>, tensor<40960xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_013_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(2) : index + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %residual_output = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<40960xi8>, tensor<40960xi8>, tensor<40960xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_014_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<139264xi8> + %second_output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_015_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<139264xi8> + %rhs = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<139264xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<139264xi8>, tensor<139264xi8>, tensor<139264xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_016_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<139264xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<17408xi8> + %scales = check.generate.fill value(0) : tensor<4352xi8> + %sums = check.generate.fill value(0) : tensor<4352xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<139264xi8>, tensor<40960xi8>, tensor<17408xi8>, tensor<4352xi8>, tensor<4352xi8>, tensor<40960xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_017_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<40960xi8>, tensor<98304xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_018_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_019_ggml_mul_mat_f32_f32_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<4300800xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<40960xi8>, tensor<4300800xi8>, tensor<8192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_019_ggml_mul_mat_f32_f32_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_019_ggml_mul_mat_f32_f32_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + %quantized_values = check.generate.fill value(0) : tensor<5120xi8> + %scales = check.generate.fill value(0) : tensor<1280xi8> + %sums = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<40960xi8>, tensor<8192xi8>, tensor<5120xi8>, tensor<1280xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_020_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_021_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<97280xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<97280xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<49152xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_021_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_021_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_022_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<8192xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<32xi8> + %output = check.generate.fill value(0) : tensor<8192xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<8192xi8>, tensor<1024xi8>, tensor<32xi8>, tensor<8192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_022_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_022_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_023_ggml_set_rows_case { + %token_count = check.literal value(2) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<2x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<2xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<2x1024xf32>, tensor<2xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_023_ggml_set_rows_case> @qwen38_27b_udq4kxl_companion_tg_c5_023_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_024_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(2) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<2x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<2x256xf16> + %output = check.generate.fill value(1.0) : tensor<2x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<2x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<2x256xf16>, tensor<2x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_024_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_024_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<6144xi8> + %scales = check.generate.fill value(0) : tensor<1536xi8> + %sums = check.generate.fill value(0) : tensor<1536xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<49152xi8>, tensor<6144xi8>, tensor<1536xi8>, tensor<1536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_025_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_026_ggml_binary_f32_case { + %element_count = check.literal value(10240) : index + %lhs = check.generate.fill value(2.0) : tensor<10240xf32> + %rhs = check.generate.fill value(3.0) : tensor<10240xf32> + %output = check.generate.fill value(0.0) : tensor<10240xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<10240xf32>, tensor<10240xf32>, tensor<10240xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_026_ggml_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_026_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_027_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(2.0) : tensor<2x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<2x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<2x5120xf32>, tensor<5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_027_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_027_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_028_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<2x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_028_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_028_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<20480xi8>, tensor<1042944000xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_029_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_030_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<40960xi8> + %rhs = check.generate.fill value(0) : tensor<40960xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<40960xi8>, tensor<40960xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_030_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_030_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_031_ggml_mul_mat_f32_f32_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<43008000xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_mul_mat_f32_f32_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<43008000xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_031_ggml_mul_mat_f32_f32_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_031_ggml_mul_mat_f32_f32_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_032_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<16xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<16xi8>, tensor<715161600xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_032_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_032_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_033_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + %partial = check.generate.fill value(0) : tensor<163840xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>, tensor<163840xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_034_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<768xi8> + %second_output = check.generate.fill value(0) : tensor<768xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<768xi8>, tensor<768xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_035_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<81920xi8>, tensor<98304xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_036_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_037_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<163840xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_038_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<131072xi8> + %k = check.generate.fill value(0) : tensor<131072xi8> + %v = check.generate.fill value(0) : tensor<147456xi8> + %g = check.generate.fill value(0) : tensor<768xi8> + %beta = check.generate.fill value(0) : tensor<768xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<40894464xi8> + %dst = check.generate.fill value(0) : tensor<98304xi8> + %rms_scales = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<131072xi8>, tensor<131072xi8>, tensor<147456xi8>, tensor<768xi8>, tensor<768xi8>, tensor<3145728xi8>, tensor<40894464xi8>, tensor<98304xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_companion_tg_c5_039_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(192) : index + %input = check.generate.fill value(0) : tensor<98304xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<98304xi8>, tensor<512xi8>, tensor<98304xi8>, tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_040_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<98304xi8>, tensor<81920xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_041_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %residual_output = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<81920xi8>, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_042_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<278528xi8> + %second_output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_043_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<278528xi8> + %rhs = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<278528xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_044_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<278528xi8>, tensor<81920xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>, tensor<81920xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_045_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<196608xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<81920xi8>, tensor<196608xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_046_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<5611520xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_047_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<81920xi8>, tensor<16384xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_048_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_049_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<195584xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<195584xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<98304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_049_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_049_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_050_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<16384xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<16384xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_050_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_050_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_051_ggml_set_rows_case { + %token_count = check.literal value(4) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<4x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<4xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<4x1024xf32>, tensor<4xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_051_ggml_set_rows_case> @qwen38_27b_udq4kxl_companion_tg_c5_051_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_052_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<4x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<4x256xf16> + %output = check.generate.fill value(1.0) : tensor<4x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<4x256xf16>, tensor<4x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_052_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_052_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_053_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_054_ggml_binary_f32_case { + %element_count = check.literal value(20480) : index + %lhs = check.generate.fill value(2.0) : tensor<20480xf32> + %rhs = check.generate.fill value(3.0) : tensor<20480xf32> + %output = check.generate.fill value(0.0) : tensor<20480xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<20480xf32>, tensor<20480xf32>, tensor<20480xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_054_ggml_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_054_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_055_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(2.0) : tensor<4x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<4x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<4x5120xf32>, tensor<5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_055_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_055_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_056_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<4xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<4x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi32>, tensor<4x5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_056_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_056_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<3973120xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<1042944000xi8>, tensor<3973120xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_057_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_058_ggml_get_rows_f32_case { + %token_count = check.literal value(5) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<20xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<20xi8>, tensor<715161600xi8>, tensor<102400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_058_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_058_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_059_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<102400xi8>, tensor<20480xi8>, tensor<102400xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_059_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_059_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + %partial = check.generate.fill value(0) : tensor<204800xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<102400xi8>, tensor<204800xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>, tensor<204800xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_060_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<960xi8> + %second_output = check.generate.fill value(0) : tensor<960xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<960xi8>, tensor<960xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_061_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_062_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<102400xi8>, tensor<122880xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_062_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_062_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_063_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<204800xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<204800xi8>, tensor<163840xi8>, tensor<204800xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_064_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_064_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_065_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<172032xi8> + %k = check.generate.fill value(0) : tensor<172032xi8> + %v = check.generate.fill value(0) : tensor<188416xi8> + %g = check.generate.fill value(0) : tensor<960xi8> + %beta = check.generate.fill value(0) : tensor<960xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<53477376xi8> + %dst = check.generate.fill value(0) : tensor<122880xi8> + %rms_scales = check.generate.fill value(0) : tensor<960xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<172032xi8>, tensor<172032xi8>, tensor<188416xi8>, tensor<960xi8>, tensor<960xi8>, tensor<3145728xi8>, tensor<53477376xi8>, tensor<122880xi8>, tensor<960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_065_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_companion_tg_c5_065_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(240) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<122880xi8>, tensor<512xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_066_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_067_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<122880xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<122880xi8>, tensor<102400xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_067_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_067_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_068_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(5) : index + %lhs = check.generate.fill value(0) : tensor<102400xi8> + %rhs = check.generate.fill value(0) : tensor<102400xi8> + %residual_output = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<102400xi8>, tensor<102400xi8>, tensor<102400xi8>, tensor<20480xi8>, tensor<102400xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_068_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_068_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<348160xi8> + %second_output = check.generate.fill value(0) : tensor<348160xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<348160xi8>, tensor<348160xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_069_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_070_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<348160xi8> + %rhs = check.generate.fill value(0) : tensor<348160xi8> + %output = check.generate.fill value(0) : tensor<348160xi8> + %quantized_values = check.generate.fill value(0) : tensor<43520xi8> + %scales = check.generate.fill value(0) : tensor<10880xi8> + %sums = check.generate.fill value(0) : tensor<10880xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<348160xi8>, tensor<348160xi8>, tensor<348160xi8>, tensor<43520xi8>, tensor<10880xi8>, tensor<10880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_070_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_070_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<348160xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + %quantized_values = check.generate.fill value(0) : tensor<43520xi8> + %scales = check.generate.fill value(0) : tensor<10880xi8> + %sums = check.generate.fill value(0) : tensor<10880xi8> + %partial = check.generate.fill value(0) : tensor<102400xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<348160xi8>, tensor<102400xi8>, tensor<43520xi8>, tensor<10880xi8>, tensor<10880xi8>, tensor<102400xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_071_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_072_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<245760xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<102400xi8>, tensor<245760xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_072_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_072_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<102400xi8>, tensor<5611520xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_073_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_074_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<12800xi8> + %scales = check.generate.fill value(0) : tensor<3200xi8> + %sums = check.generate.fill value(0) : tensor<3200xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<102400xi8>, tensor<20480xi8>, tensor<12800xi8>, tensor<3200xi8>, tensor<3200xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_074_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_074_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_075_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<244736xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<80xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<244736xi8>, tensor<1024xi8>, tensor<80xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_075_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_075_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_076_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<80xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<20480xi8>, tensor<1024xi8>, tensor<80xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_076_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_076_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_077_ggml_set_rows_case { + %token_count = check.literal value(5) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<5x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<5xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<5x1024xf32>, tensor<5xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_077_ggml_set_rows_case> @qwen38_27b_udq4kxl_companion_tg_c5_077_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_078_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(5) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<5x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<5x256xf16> + %output = check.generate.fill value(1.0) : tensor<5x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<5x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<5x256xf16>, tensor<5x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_078_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_078_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_079_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<15360xi8> + %scales = check.generate.fill value(0) : tensor<3840xi8> + %sums = check.generate.fill value(0) : tensor<3840xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<122880xi8>, tensor<15360xi8>, tensor<3840xi8>, tensor<3840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_079_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_079_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_080_ggml_binary_f32_case { + %element_count = check.literal value(25600) : index + %lhs = check.generate.fill value(2.0) : tensor<25600xf32> + %rhs = check.generate.fill value(3.0) : tensor<25600xf32> + %output = check.generate.fill value(0.0) : tensor<25600xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<25600xf32>, tensor<25600xf32>, tensor<25600xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_080_ggml_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_080_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_081_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(2.0) : tensor<5x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<5x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<5x5120xf32>, tensor<5120xf32>, tensor<5x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_081_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_081_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_082_ggml_get_rows_f32_case { + %token_count = check.literal value(5) : index + %row_count = check.literal value(5) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<5xi32> + %weight = check.generate.fill value(0.0) : tensor<5x5120xf32> + %output = check.generate.fill value(1.0) : tensor<5x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<5xi32>, tensor<5x5120xf32>, tensor<5x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_082_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_082_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<102400xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<4966400xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<102400xi8>, tensor<1042944000xi8>, tensor<4966400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_083_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_084_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<102400xi8> + %rhs = check.generate.fill value(0) : tensor<102400xi8> + %output = check.generate.fill value(0) : tensor<204800xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<102400xi8>, tensor<102400xi8>, tensor<204800xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_084_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_084_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(5) : index + %input = check.generate.fill value(0) : tensor<204800xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<102400xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<204800xi8>, tensor<56115200xi8>, tensor<102400xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_085_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_086_ggml_get_rows_f32_case { + %token_count = check.literal value(3) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<12xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<12xi8>, tensor<715161600xi8>, tensor<61440xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_086_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_086_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_087_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<61440xi8>, tensor<20480xi8>, tensor<61440xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_087_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_087_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + %partial = check.generate.fill value(0) : tensor<122880xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<61440xi8>, tensor<122880xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>, tensor<122880xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_088_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<576xi8> + %second_output = check.generate.fill value(0) : tensor<576xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<576xi8>, tensor<576xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_089_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_090_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<73728xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<61440xi8>, tensor<73728xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_090_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_090_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_091_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<122880xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<122880xi8>, tensor<163840xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_092_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_092_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_093_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<90112xi8> + %k = check.generate.fill value(0) : tensor<90112xi8> + %v = check.generate.fill value(0) : tensor<106496xi8> + %g = check.generate.fill value(0) : tensor<576xi8> + %beta = check.generate.fill value(0) : tensor<576xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<28311552xi8> + %dst = check.generate.fill value(0) : tensor<73728xi8> + %rms_scales = check.generate.fill value(0) : tensor<576xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<90112xi8>, tensor<90112xi8>, tensor<106496xi8>, tensor<576xi8>, tensor<576xi8>, tensor<3145728xi8>, tensor<28311552xi8>, tensor<73728xi8>, tensor<576xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_093_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @qwen38_27b_udq4kxl_companion_tg_c5_093_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(144) : index + %input = check.generate.fill value(0) : tensor<73728xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<73728xi8> + %output = check.generate.fill value(0) : tensor<73728xi8> + %quantized_values = check.generate.fill value(0) : tensor<9216xi8> + %scales = check.generate.fill value(0) : tensor<2304xi8> + %sums = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<73728xi8>, tensor<512xi8>, tensor<73728xi8>, tensor<73728xi8>, tensor<9216xi8>, tensor<2304xi8>, tensor<2304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_094_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_095_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<73728xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<9216xi8> + %scales = check.generate.fill value(0) : tensor<2304xi8> + %sums = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<73728xi8>, tensor<61440xi8>, tensor<9216xi8>, tensor<2304xi8>, tensor<2304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_095_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_095_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_096_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(3) : index + %lhs = check.generate.fill value(0) : tensor<61440xi8> + %rhs = check.generate.fill value(0) : tensor<61440xi8> + %residual_output = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<61440xi8>, tensor<61440xi8>, tensor<61440xi8>, tensor<20480xi8>, tensor<61440xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_096_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_096_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<208896xi8> + %second_output = check.generate.fill value(0) : tensor<208896xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<208896xi8>, tensor<208896xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_097_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_098_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<208896xi8> + %rhs = check.generate.fill value(0) : tensor<208896xi8> + %output = check.generate.fill value(0) : tensor<208896xi8> + %quantized_values = check.generate.fill value(0) : tensor<26112xi8> + %scales = check.generate.fill value(0) : tensor<6528xi8> + %sums = check.generate.fill value(0) : tensor<6528xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<208896xi8>, tensor<208896xi8>, tensor<208896xi8>, tensor<26112xi8>, tensor<6528xi8>, tensor<6528xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_098_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_098_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<208896xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + %quantized_values = check.generate.fill value(0) : tensor<26112xi8> + %scales = check.generate.fill value(0) : tensor<6528xi8> + %sums = check.generate.fill value(0) : tensor<6528xi8> + %partial = check.generate.fill value(0) : tensor<61440xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<208896xi8>, tensor<61440xi8>, tensor<26112xi8>, tensor<6528xi8>, tensor<6528xi8>, tensor<61440xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_099_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_100_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<147456xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<61440xi8>, tensor<147456xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_100_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_100_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<61440xi8>, tensor<5611520xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_101_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_102_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<12288xi8> + %quantized_values = check.generate.fill value(0) : tensor<7680xi8> + %scales = check.generate.fill value(0) : tensor<1920xi8> + %sums = check.generate.fill value(0) : tensor<1920xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<61440xi8>, tensor<12288xi8>, tensor<7680xi8>, tensor<1920xi8>, tensor<1920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_102_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_102_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_103_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<146432xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<48xi8> + %output = check.generate.fill value(0) : tensor<73728xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<146432xi8>, tensor<1024xi8>, tensor<48xi8>, tensor<73728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_103_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_103_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_104_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<12288xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<48xi8> + %output = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<12288xi8>, tensor<1024xi8>, tensor<48xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_104_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_104_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_105_ggml_set_rows_case { + %token_count = check.literal value(3) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<3x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<3xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<3x1024xf32>, tensor<3xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_105_ggml_set_rows_case> @qwen38_27b_udq4kxl_companion_tg_c5_105_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_106_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(3) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<3x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<3x256xf16> + %output = check.generate.fill value(1.0) : tensor<3x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<3x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<3x256xf16>, tensor<3x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_106_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_106_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_107_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<73728xi8> + %quantized_values = check.generate.fill value(0) : tensor<9216xi8> + %scales = check.generate.fill value(0) : tensor<2304xi8> + %sums = check.generate.fill value(0) : tensor<2304xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<73728xi8>, tensor<9216xi8>, tensor<2304xi8>, tensor<2304xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_107_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_107_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_108_ggml_binary_f32_case { + %element_count = check.literal value(15360) : index + %lhs = check.generate.fill value(2.0) : tensor<15360xf32> + %rhs = check.generate.fill value(3.0) : tensor<15360xf32> + %output = check.generate.fill value(0.0) : tensor<15360xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<15360xf32>, tensor<15360xf32>, tensor<15360xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_108_ggml_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_108_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_109_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(2.0) : tensor<3x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<3x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<3x5120xf32>, tensor<5120xf32>, tensor<3x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_109_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_109_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_110_ggml_get_rows_f32_case { + %token_count = check.literal value(3) : index + %row_count = check.literal value(3) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<3xi32> + %weight = check.generate.fill value(0.0) : tensor<3x5120xf32> + %output = check.generate.fill value(1.0) : tensor<3x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<3xi32>, tensor<3x5120xf32>, tensor<3x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_110_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_110_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<61440xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<2979840xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<61440xi8>, tensor<1042944000xi8>, tensor<2979840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_111_ggml_mul_mat_q6_k_packed_token1_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_112_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<61440xi8> + %rhs = check.generate.fill value(0) : tensor<61440xi8> + %output = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<61440xi8>, tensor<61440xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_112_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_112_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(3) : index + %input = check.generate.fill value(0) : tensor<122880xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<61440xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<122880xi8>, tensor<56115200xi8>, tensor<61440xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_113_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_114_ggml_get_rows_f32_case { + %token_count = check.literal value(16) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<64xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<64xi8>, tensor<715161600xi8>, tensor<327680xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_114_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_114_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_115_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<327680xi8>, tensor<20480xi8>, tensor<327680xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_115_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_115_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + %partial = check.generate.fill value(0) : tensor<655360xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<327680xi8>, tensor<655360xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>, tensor<655360xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_116_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<3072xi8> + %second_output = check.generate.fill value(0) : tensor<3072xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<3072xi8>, tensor<3072xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_117_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_118_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<327680xi8>, tensor<393216xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_118_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_118_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_119_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<655360xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<655360xi8>, tensor<163840xi8>, tensor<655360xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_120_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_120_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_121_llm_gated_delta_net_f32_wmma_head128_case { + %q = check.generate.fill value(0) : tensor<417792xi8> + %k = check.generate.fill value(0) : tensor<417792xi8> + %v = check.generate.fill value(0) : tensor<434176xi8> + %g = check.generate.fill value(0) : tensor<2112xi8> + %beta = check.generate.fill value(0) : tensor<2112xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<3416064xi8> + %rms_scales = check.generate.fill value(0) : tensor<2112xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128[](%q, %k, %v, %g, %beta, %state_in, %dst, %rms_scales) : [](tensor<417792xi8>, tensor<417792xi8>, tensor<434176xi8>, tensor<2112xi8>, tensor<2112xi8>, tensor<3145728xi8>, tensor<3416064xi8>, tensor<2112xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_121_llm_gated_delta_net_f32_wmma_head128_case> @qwen38_27b_udq4kxl_companion_tg_c5_121_llm_gated_delta_net_f32_wmma_head128 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_122_ggml_copy_f32_case { + %element_count = check.literal value(786432) : index + %source = check.generate.fill value(0) : tensor<3145728xi8> + %output = check.generate.fill value(0) : tensor<3145728xi8> + kernel.launch @ggml_copy_f32[%element_count](%element_count, %source, %output) : [index](index, tensor<3145728xi8>, tensor<3145728xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_122_ggml_copy_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_122_ggml_copy_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_123_llm_gated_delta_net_f32_wmma_head128_inplace_case { + %q = check.generate.fill value(0) : tensor<8192xi8> + %k = check.generate.fill value(0) : tensor<8192xi8> + %v = check.generate.fill value(0) : tensor<24576xi8> + %g = check.generate.fill value(0) : tensor<192xi8> + %beta = check.generate.fill value(0) : tensor<192xi8> + %state_inout = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<24576xi8> + %rms_scales = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_inplace[](%q, %k, %v, %g, %beta, %state_inout, %dst, %rms_scales) : [](tensor<8192xi8>, tensor<8192xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3145728xi8>, tensor<24576xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_123_llm_gated_delta_net_f32_wmma_head128_inplace_case> @qwen38_27b_udq4kxl_companion_tg_c5_123_llm_gated_delta_net_f32_wmma_head128_inplace + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(768) : index + %input = check.generate.fill value(0) : tensor<393216xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<393216xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<393216xi8>, tensor<512xi8>, tensor<393216xi8>, tensor<393216xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_124_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_125_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<393216xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<393216xi8>, tensor<327680xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_125_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_125_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_126_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(16) : index + %lhs = check.generate.fill value(0) : tensor<327680xi8> + %rhs = check.generate.fill value(0) : tensor<327680xi8> + %residual_output = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<327680xi8>, tensor<327680xi8>, tensor<327680xi8>, tensor<20480xi8>, tensor<327680xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_126_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_126_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<1114112xi8> + %second_output = check.generate.fill value(0) : tensor<1114112xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<1114112xi8>, tensor<1114112xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_127_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_128_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<1114112xi8> + %rhs = check.generate.fill value(0) : tensor<1114112xi8> + %output = check.generate.fill value(0) : tensor<1114112xi8> + %quantized_values = check.generate.fill value(0) : tensor<139264xi8> + %scales = check.generate.fill value(0) : tensor<34816xi8> + %sums = check.generate.fill value(0) : tensor<34816xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<1114112xi8>, tensor<1114112xi8>, tensor<1114112xi8>, tensor<139264xi8>, tensor<34816xi8>, tensor<34816xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_128_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_128_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<1114112xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + %quantized_values = check.generate.fill value(0) : tensor<139264xi8> + %scales = check.generate.fill value(0) : tensor<34816xi8> + %sums = check.generate.fill value(0) : tensor<34816xi8> + %partial = check.generate.fill value(0) : tensor<327680xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<1114112xi8>, tensor<327680xi8>, tensor<139264xi8>, tensor<34816xi8>, tensor<34816xi8>, tensor<327680xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_129_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<786432xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<327680xi8>, tensor<786432xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_130_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<327680xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<327680xi8>, tensor<5611520xi8>, tensor<65536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_131_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + %quantized_values = check.generate.fill value(0) : tensor<40960xi8> + %scales = check.generate.fill value(0) : tensor<10240xi8> + %sums = check.generate.fill value(0) : tensor<10240xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<327680xi8>, tensor<65536xi8>, tensor<40960xi8>, tensor<10240xi8>, tensor<10240xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_132_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_133_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<785408xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<393216xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<785408xi8>, tensor<1024xi8>, tensor<256xi8>, tensor<393216xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_133_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_133_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_134_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<65536xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<65536xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<65536xi8>, tensor<1024xi8>, tensor<256xi8>, tensor<65536xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_134_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_134_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_135_ggml_set_rows_case { + %token_count = check.literal value(16) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<16x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<16xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<16x1024xf32>, tensor<16xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_135_ggml_set_rows_case> @qwen38_27b_udq4kxl_companion_tg_c5_135_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_136_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(16) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<16x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<16x256xf16> + %output = check.generate.fill value(1.0) : tensor<16x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<16x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<16x256xf16>, tensor<16x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_136_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_136_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<393216xi8> + %quantized_values = check.generate.fill value(0) : tensor<49152xi8> + %scales = check.generate.fill value(0) : tensor<12288xi8> + %sums = check.generate.fill value(0) : tensor<12288xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<393216xi8>, tensor<49152xi8>, tensor<12288xi8>, tensor<12288xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_137_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_138_ggml_binary_f32_case { + %element_count = check.literal value(81920) : index + %lhs = check.generate.fill value(2.0) : tensor<81920xf32> + %rhs = check.generate.fill value(3.0) : tensor<81920xf32> + %output = check.generate.fill value(0.0) : tensor<81920xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<81920xf32>, tensor<81920xf32>, tensor<81920xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_138_ggml_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_138_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_139_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(2.0) : tensor<16x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<16x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<16x5120xf32>, tensor<5120xf32>, tensor<16x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_139_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_139_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_140_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<327680xi8> + %rhs = check.generate.fill value(0) : tensor<327680xi8> + %output = check.generate.fill value(0) : tensor<655360xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<327680xi8>, tensor<327680xi8>, tensor<655360xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_140_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_140_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(16) : index + %input = check.generate.fill value(0) : tensor<655360xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<327680xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<655360xi8>, tensor<56115200xi8>, tensor<327680xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_141_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_142_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<4x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_142_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_142_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_143_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<81920xi8>, tensor<81920xi8>, tensor<163840xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_143_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_143_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<163840xi8> + %weight = check.generate.fill value(0) : tensor<56115200xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<163840xi8>, tensor<56115200xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_144_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_145_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<4xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<4xi8>, tensor<715161600xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_145_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_145_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_146_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(2.0) : tensor<1x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<1x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<1x5120xf32>, tensor<5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_146_ggml_rmsnorm_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_146_ggml_rmsnorm_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_147_ggml_concat_dim0_f32_case { + %lhs = check.generate.fill value(0) : tensor<20480xi8> + %rhs = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + kernel.launch @ggml_concat_dim0_f32[](%lhs, %rhs, %output) : [](tensor<20480xi8>, tensor<20480xi8>, tensor<40960xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_147_ggml_concat_dim0_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_147_ggml_concat_dim0_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_148_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(10240) : index + %output_size = check.literal value(5120) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<43008000xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<40960xi8>, tensor<43008000xi8>, tensor<20480xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_148_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_companion_tg_c5_148_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_149_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_149_ggml_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_149_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_150_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<49152xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<20480xi8>, tensor<49152xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_150_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_150_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_151_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(5120) : index + %output_size = check.literal value(1024) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<4300800xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<20480xi8>, tensor<4300800xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_151_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_companion_tg_c5_151_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_152_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<20480xi8>, tensor<4096xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_152_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_152_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_153_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<48128xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<16xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<48128xi8>, tensor<1024xi8>, tensor<16xi8>, tensor<24576xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_153_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_153_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_154_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<4096xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<16xi8> + %output = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<4096xi8>, tensor<1024xi8>, tensor<16xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_154_ggml_rmsnorm_mul_rope_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_154_ggml_rmsnorm_mul_rope_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_155_ggml_set_rows_case { + %token_count = check.literal value(1) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<1x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<1xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<1x1024xf32>, tensor<1xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_155_ggml_set_rows_case> @qwen38_27b_udq4kxl_companion_tg_c5_155_ggml_set_rows + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_156_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(1) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<1x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<1x256xf16> + %output = check.generate.fill value(1.0) : tensor<1x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<1x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<1x256xf16>, tensor<1x24x256xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_156_ggml_flash_attention_f32_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_156_ggml_flash_attention_f32_f16_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_157_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<24576xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_157_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_157_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_158_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<24576xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<24576xi8>, tensor<20480xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_158_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_158_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_159_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(1) : index + %lhs = check.generate.fill value(0) : tensor<20480xi8> + %rhs = check.generate.fill value(0) : tensor<20480xi8> + %residual_output = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_159_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_159_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<69632xi8> + %second_output = check.generate.fill value(0) : tensor<69632xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<69632xi8>, tensor<69632xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_160_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_161_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<69632xi8> + %rhs = check.generate.fill value(0) : tensor<69632xi8> + %output = check.generate.fill value(0) : tensor<69632xi8> + %quantized_values = check.generate.fill value(0) : tensor<8704xi8> + %scales = check.generate.fill value(0) : tensor<2176xi8> + %sums = check.generate.fill value(0) : tensor<2176xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<69632xi8>, tensor<69632xi8>, tensor<69632xi8>, tensor<8704xi8>, tensor<2176xi8>, tensor<2176xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_161_ggml_binary_swiglu_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_161_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<69632xi8> + %output = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<8704xi8> + %scales = check.generate.fill value(0) : tensor<2176xi8> + %sums = check.generate.fill value(0) : tensor<2176xi8> + %partial = check.generate.fill value(0) : tensor<20480xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<69632xi8>, tensor<20480xi8>, tensor<8704xi8>, tensor<2176xi8>, tensor<2176xi8>, tensor<20480xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_162_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_163_ggml_binary_f32_case { + %element_count = check.literal value(5120) : index + %lhs = check.generate.fill value(2.0) : tensor<5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<5120xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<5120xf32>, tensor<5120xf32>, tensor<5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_163_ggml_binary_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_163_ggml_binary_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_164_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(1) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<1x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<1x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_164_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_164_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_165_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<20480xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_165_ggml_quantize_f32_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_165_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_166_ggml_select_symmetric_i4_k32_groups_case { + %input = check.generate.fill value(0) : tensor<20480xi8> + %selected_mask = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_select_symmetric_i4_k32_groups[](%input, %selected_mask) : [](tensor<20480xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_166_ggml_select_symmetric_i4_k32_groups_case> @qwen38_27b_udq4kxl_companion_tg_c5_166_ggml_select_symmetric_i4_k32_groups + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_167_ggml_mul_mat_q6_k_symmetric_i2_scan_token1_case { + %weight = check.generate.fill value(0) : tensor<1380659200xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + %qact = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %selected_mask = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_q6_k_symmetric_i2_scan_token1[](%weight, %output, %qact, %scales, %selected_mask) : [](tensor<1380659200xi8>, tensor<993280xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_167_ggml_mul_mat_q6_k_symmetric_i2_scan_token1_case> @qwen38_27b_udq4kxl_companion_tg_c5_167_ggml_mul_mat_q6_k_symmetric_i2_scan_token1 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_168_ggml_top_k8_f32_partitions_register_case { + %element_count = check.literal value(248320) : index + %values = check.generate.fill value(0) : tensor<993280xi8> + %partial_values = check.generate.fill value(0) : tensor<4096xi8> + %partial_ids = check.generate.fill value(0) : tensor<4096xi8> + kernel.launch @ggml_top_k8_f32_partitions_register[%element_count](%element_count, %values, %partial_values, %partial_ids) : [index](index, tensor<993280xi8>, tensor<4096xi8>, tensor<4096xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_168_ggml_top_k8_f32_partitions_register_case> @qwen38_27b_udq4kxl_companion_tg_c5_168_ggml_top_k8_f32_partitions_register + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_169_ggml_top_k128_f32_reduce_gather_register_case { + %element_count = check.literal value(248320) : index + %partial_values = check.generate.fill value(0) : tensor<4096xi8> + %partial_ids = check.generate.fill value(0) : tensor<4096xi8> + %candidate_output = check.generate.fill value(0) : tensor<512xi8> + %value_output = check.generate.fill value(0) : tensor<512xi8> + kernel.launch @ggml_top_k128_f32_reduce_gather_register[%element_count](%element_count, %partial_values, %partial_ids, %candidate_output, %value_output) : [index](index, tensor<4096xi8>, tensor<4096xi8>, tensor<512xi8>, tensor<512xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_169_ggml_top_k128_f32_reduce_gather_register_case> @qwen38_27b_udq4kxl_companion_tg_c5_169_ggml_top_k128_f32_reduce_gather_register + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_170_ggml_fill_negative_f32_case { + %element_count = check.literal value(248320) : index + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_fill_negative_f32[%element_count](%element_count, %output) : [index](index, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_170_ggml_fill_negative_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_170_ggml_fill_negative_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_171_ggml_mul_mat_q6_k_packed_selected_refine_token1_case { + %token_count = check.literal value(1) : index + %candidate_count = check.literal value(64) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1380659200xi8> + %candidates = check.generate.fill value(0) : tensor<256xi8> + %exact_output = check.generate.fill value(0) : tensor<256xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_selected_refine_token1[%token_count, %candidate_count](%token_count, %candidate_count, %input, %weight, %candidates, %exact_output, %output) : [index, index](index, index, tensor<20480xi8>, tensor<1380659200xi8>, tensor<256xi8>, tensor<256xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_171_ggml_mul_mat_q6_k_packed_selected_refine_token1_case> @qwen38_27b_udq4kxl_companion_tg_c5_171_ggml_mul_mat_q6_k_packed_selected_refine_token1 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_172_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + %partial = check.generate.fill value(0) : tensor<40960xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<20480xi8>, tensor<40960xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>, tensor<40960xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_172_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_172_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_173_ggml_mul_mat_f32_f32_decode_wave64_case { + %token_count = check.literal value(1) : index + %input_size = check.literal value(5120) : index + %output_size = check.literal value(48) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<138240xi8> + %output = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @ggml_mul_mat_f32_f32_decode_wave64[%token_count, %input_size, %output_size](%token_count, %input_size, %output_size, %input, %weight, %output) : [index, index, index](index, index, index, tensor<20480xi8>, tensor<138240xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_173_ggml_mul_mat_f32_f32_decode_wave64_case> @qwen38_27b_udq4kxl_companion_tg_c5_173_ggml_mul_mat_f32_f32_decode_wave64 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_174_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<2560xi8> + %scales = check.generate.fill value(0) : tensor<640xi8> + %sums = check.generate.fill value(0) : tensor<640xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<20480xi8>, tensor<24576xi8>, tensor<2560xi8>, tensor<640xi8>, tensor<640xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_174_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_174_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_175_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<40960xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<40960xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<40960xi8>, tensor<163840xi8>, tensor<40960xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +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[%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 +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_176_llm_gated_delta_net_projection_epilogue_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_176_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_177_llm_gated_delta_net_f32_wmma_head128_inplace_case { + %q = check.generate.fill value(0) : tensor<8192xi8> + %k = check.generate.fill value(0) : tensor<8192xi8> + %v = check.generate.fill value(0) : tensor<24576xi8> + %g = check.generate.fill value(0) : tensor<192xi8> + %beta = check.generate.fill value(0) : tensor<192xi8> + %state_inout = check.generate.fill value(0) : tensor<3145728xi8> + %dst = check.generate.fill value(0) : tensor<24576xi8> + %rms_scales = check.generate.fill value(0) : tensor<192xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_inplace[](%q, %k, %v, %g, %beta, %state_inout, %dst, %rms_scales) : [](tensor<8192xi8>, tensor<8192xi8>, tensor<24576xi8>, tensor<192xi8>, tensor<192xi8>, tensor<3145728xi8>, tensor<24576xi8>, tensor<192xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_177_llm_gated_delta_net_f32_wmma_head128_inplace_case> @qwen38_27b_udq4kxl_companion_tg_c5_177_llm_gated_delta_net_f32_wmma_head128_inplace + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_178_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(48) : index + %input = check.generate.fill value(0) : tensor<24576xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<24576xi8> + %output = check.generate.fill value(0) : tensor<24576xi8> + %quantized_values = check.generate.fill value(0) : tensor<3072xi8> + %scales = check.generate.fill value(0) : tensor<768xi8> + %sums = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<24576xi8>, tensor<512xi8>, tensor<24576xi8>, tensor<24576xi8>, tensor<3072xi8>, tensor<768xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_178_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @qwen38_27b_udq4kxl_companion_tg_c5_178_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_179_ggml_get_rows_f32_case { + %token_count = check.literal value(2) : index + %row_count = check.literal value(2) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<2xi32> + %weight = check.generate.fill value(0.0) : tensor<2x5120xf32> + %output = check.generate.fill value(1.0) : tensor<2x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<2xi32>, tensor<2x5120xf32>, tensor<2x5120xf32>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_179_ggml_get_rows_f32_case> @qwen38_27b_udq4kxl_companion_tg_c5_179_ggml_get_rows_f32 + +check.case public @qwen38_27b_udq4kxl_companion_tg_c5_180_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(2) : index + %input = check.generate.fill value(0) : tensor<40960xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<1986560xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<40960xi8>, tensor<1042944000xi8>, tensor<1986560xi8>) + check.return +} + +check.benchmark<@qwen38_27b_udq4kxl_companion_tg_c5_180_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @qwen38_27b_udq4kxl_companion_tg_c5_180_ggml_mul_mat_q6_k_packed_token1_f16_wmma 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 new file mode 100644 index 000000000000..5fb4222c0d2b --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.tg_c4.json @@ -0,0 +1,965 @@ +{ + "command_count": 965, + "dispatch_count": 28, + "dispatches": [ + { + "benchmark": "@v2_mtp4_program4_tg_c4_000_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "4", + "ggml.get_rows_f32.weight_format": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 248320, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_001_ggml_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "10240", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "48", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_004_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_005_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "30720", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "hidden_size": 30720, + "row_count": 20, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_006_llm_ssm_conv_dconv4_silu_rollback_f32", + "compile_parameters": { + "llm.ssm_conv.rollback.cache_count": "5", + "llm.ssm_conv.rollback.d_inner": "10240", + "llm.ssm_conv.rollback.n_s": "1", + "llm.ssm_conv.rollback.n_t": "4", + "llm.ssm_conv.rollback.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32", + "library_sources": [ + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "primary_sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom" + ], + "sources": [ + "ops/ssm_conv/rollback_dconv4_silu_f32.loom", + "motifs/ssm_conv/state_materialize.loom", + "motifs/ssm_conv/rollback_store.loom", + "motifs/ssm_conv/decode_update.loom", + "motifs/ssm_conv/dconv4_silu.loom" + ], + "symbol": "llm_ssm_conv_dconv4_silu_rollback_f32", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_007_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "786432", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "hidden_size": 786432, + "row_count": 20, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_008_llm_gated_delta_net_projection_epilogue_f32", + "compile_parameters": { + "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": { + "element_count": 192 + }, + "kernel": "loom_libs:llm_gated_delta_net_projection_epilogue_f32", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_projection_epilogue_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_009_llm_gated_delta_net_f32_wmma_head128_snapshot", + "compile_parameters": { + "llm.gated_delta_net.head_count": "48", + "llm.gated_delta_net.head_width": "128", + "llm.gated_delta_net.l2_epsilon": "9.99999997e-07", + "llm.gated_delta_net.qk_stride1": "128", + "llm.gated_delta_net.qk_stride2": "10240", + "llm.gated_delta_net.qk_stride3": "40960", + "llm.gated_delta_net.query_head_count": "16", + "llm.gated_delta_net.query_sequence_ratio": "1", + "llm.gated_delta_net.scalar_stride1": "1", + "llm.gated_delta_net.scalar_stride2": "48", + "llm.gated_delta_net.scalar_stride3": "192", + "llm.gated_delta_net.sequence_count": "1", + "llm.gated_delta_net.snapshot_stride": "3145728", + "llm.gated_delta_net.token_count": "4", + "llm.gated_delta_net.value_stride1": "128", + "llm.gated_delta_net.value_stride2": "10240", + "llm.gated_delta_net.value_stride3": "40960", + "llm.gated_delta_net.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": {}, + "kernel": "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot", + "library_sources": [], + "primary_sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "sources": [ + "ops/gated_delta_net_f32_wmma.loom" + ], + "symbol": "llm_gated_delta_net_f32_wmma_head128_snapshot", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "compile_parameters": { + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.hidden_size": "128", + "ggml.rmsnorm_gate_silu_mul_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 48, + "integer_parameters": { + "token_count": 192 + }, + "kernel": "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_gate_silu_mul_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_011_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "6144", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 64, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_012_ggml_add_rmsnorm_binary_symmetric_i4_k32", + "compile_parameters": { + "ggml.add_rmsnorm_binary_symmetric_i4.hidden_size": "5120", + "ggml.add_rmsnorm_binary_symmetric_i4.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 127, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "primary_sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom" + ], + "sources": [ + "ops/normalization/add_rmsnorm_binary_symmetric_i4_k32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_symmetric_i4.loom" + ], + "symbol": "ggml_add_rmsnorm_binary_symmetric_i4_k32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 64, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_014_ggml_binary_swiglu_symmetric_i4_k32", + "compile_parameters": { + "ggml.binary_swiglu_symmetric_i4.input_size": "17408", + "ggml.binary_swiglu_symmetric_i4.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 64, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_binary_swiglu_symmetric_i4_k32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_swiglu_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "17408", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 64, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_016_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "12288", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_017_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "1024" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_i8_prepacked_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_018_ggml_mul_mat_symmetric_i4_lowrow_wmma", + "compile_parameters": { + "ggml.mul_mat.symmetric_i4.lowrow.input_size": "5120", + "ggml.mul_mat.symmetric_i4.lowrow.output_size": "1024", + "ggml.mul_mat.symmetric_i4.lowrow.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_mul_mat_symmetric_i4_lowrow_wmma", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_019_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "512", + "ggml.rmsnorm_mul_rope.input_stride2": "12288", + "ggml.rmsnorm_mul_rope.input_stride3": "49152", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "24", + "ggml.rmsnorm_mul_rope.ne2": "4", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "6144", + "ggml.rmsnorm_mul_rope.output_stride3": "24576", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_020_ggml_rmsnorm_mul_rope_f32", + "compile_parameters": { + "ggml.rmsnorm_mul_rope.attn_factor": "1", + "ggml.rmsnorm_mul_rope.epsilon": "9.99999997e-07", + "ggml.rmsnorm_mul_rope.freq_base": "10000000", + "ggml.rmsnorm_mul_rope.freq_scale": "1", + "ggml.rmsnorm_mul_rope.hidden_size": "256", + "ggml.rmsnorm_mul_rope.input_stride1": "256", + "ggml.rmsnorm_mul_rope.input_stride2": "1024", + "ggml.rmsnorm_mul_rope.input_stride3": "4096", + "ggml.rmsnorm_mul_rope.mode": "40", + "ggml.rmsnorm_mul_rope.n_dims": "64", + "ggml.rmsnorm_mul_rope.ne1": "4", + "ggml.rmsnorm_mul_rope.ne2": "4", + "ggml.rmsnorm_mul_rope.ne3": "1", + "ggml.rmsnorm_mul_rope.output_stride1": "256", + "ggml.rmsnorm_mul_rope.output_stride2": "1024", + "ggml.rmsnorm_mul_rope.output_stride3": "4096", + "ggml.rmsnorm_mul_rope.section0": "11", + "ggml.rmsnorm_mul_rope.section1": "11", + "ggml.rmsnorm_mul_rope.section2": "10", + "ggml.rmsnorm_mul_rope.section3": "0", + "ggml.rmsnorm_mul_rope.workgroup_size": "256" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_rmsnorm_mul_rope_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_mul_rope_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/rope/rope_f32.loom" + ], + "symbol": "ggml_rmsnorm_mul_rope_f32", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_021_ggml_set_rows", + "compile_parameters": { + "ggml.set_rows.hidden_capacity": "1024", + "ggml.set_rows.input_format": "32", + "ggml.set_rows.output_format": "16", + "ggml.set_rows.token_capacity": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 32, + "integer_parameters": { + "cache_row_count": 262144, + "hidden_size": 1024, + "token_count": 4 + }, + "kernel": "loom_libs:ggml_set_rows", + "library_sources": [ + "motifs/rows/cache_publish.loom" + ], + "primary_sources": [ + "ops/rows/set_rows.loom" + ], + "sources": [ + "ops/rows/set_rows.loom", + "motifs/rows/cache_publish.loom" + ], + "symbol": "ggml_set_rows", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "cache_row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_022_ggml_flash_attention_f32_f16_wmma", + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + "ggml.flash_attention.head_size": "256", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.query_head_count": "24" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": { + "key_value_token_count": 256, + "query_token_count": 4 + }, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "library_sources": [ + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "primary_sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom" + ], + "sources": [ + "ops/flash_attention/prefill_f32_f16_wmma.loom", + "motifs/flash_attention/qk_block_wmma.loom", + "motifs/flash_attention/online_softmax.loom", + "motifs/flash_attention/pv_block_wmma.loom", + "motifs/flash_attention/publish_f32.loom" + ], + "symbol": "ggml_flash_attention_f32_f16_wmma", + "workload_parameters": [ + { + "name": "query_token_count", + "type": "index" + }, + { + "name": "key_value_token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_023_ggml_quantize_f32_symmetric_i4_k32", + "compile_parameters": { + "ggml.quantize_symmetric_i4_k32.input_size": "6144", + "ggml.quantize_symmetric_i4_k32.token_count": "4" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 16, + "integer_parameters": {}, + "kernel": "loom_libs:ggml_quantize_f32_symmetric_i4_k32", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "sources": [ + "ops/mul_mat_swiglu_symmetric_i4_wmma.loom" + ], + "symbol": "ggml_quantize_f32_symmetric_i4_k32", + "workload_parameters": [] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_024_ggml_binary_f32", + "compile_parameters": { + "ggml.binary_f32.op": "0" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "element_count": 20480 + }, + "kernel": "loom_libs:ggml_binary_f32", + "library_sources": [ + "motifs/binary_f32_apply.loom" + ], + "primary_sources": [ + "ops/binary_f32.loom" + ], + "sources": [ + "ops/binary_f32.loom", + "motifs/binary_f32_apply.loom" + ], + "symbol": "ggml_binary_f32", + "workload_parameters": [ + { + "name": "element_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_025_ggml_rmsnorm_binary_f32", + "compile_parameters": { + "ggml.rmsnorm_binary_f32.hidden_size": "5120", + "ggml.rmsnorm_binary_f32.op": "2", + "ggml.rmsnorm_binary_f32.rms_epsilon": "9.99999997e-07" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 4 + }, + "kernel": "loom_libs:ggml_rmsnorm_binary_f32", + "library_sources": [ + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "primary_sources": [ + "ops/normalization/rmsnorm_binary_f32.loom" + ], + "sources": [ + "ops/normalization/rmsnorm_binary_f32.loom", + "motifs/normalization/rmsnorm.loom", + "motifs/normalization/rmsnorm_binary.loom", + "motifs/elementwise/binary_apply.loom" + ], + "symbol": "ggml_rmsnorm_binary_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_026_ggml_get_rows_f32", + "compile_parameters": { + "ggml.get_rows_f32.hidden_capacity": "5120", + "ggml.get_rows_f32.token_capacity": "1", + "ggml.get_rows_f32.weight_format": "32" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "hidden_size": 5120, + "row_count": 4, + "token_count": 1 + }, + "kernel": "loom_libs:ggml_get_rows_f32", + "library_sources": [ + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "primary_sources": [ + "ops/rows/get_rows_f32.loom" + ], + "sources": [ + "ops/rows/get_rows_f32.loom", + "motifs/dequant.loom", + "motifs/publish_f32.loom", + "motifs/quantize_q8_1_x4.loom", + "motifs/q4_k_f16.loom", + "motifs/q6_k_f16.loom", + "motifs/q8_0_f16.loom", + "motifs/q8_1_f16.loom", + "motifs/f16_f16.loom", + "motifs/bf16_f16.loom", + "motifs/f32_f16.loom" + ], + "symbol": "ggml_get_rows_f32", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + }, + { + "name": "row_count", + "type": "index" + }, + { + "name": "hidden_size", + "type": "index" + } + ] + }, + { + "benchmark": "@v2_mtp4_program4_tg_c4_027_ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "compile_parameters": { + "ggml.mul_mat_q6_k_packed.input_size": "5120", + "ggml.mul_mat_q6_k_packed.output_accumulation": "0", + "ggml.mul_mat_q6_k_packed.output_size": "248320" + }, + "corpus_dir": "kernel-corpus/kernels/loom-libs", + "count": 1, + "integer_parameters": { + "token_count": 1 + }, + "kernel": "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "library_sources": [], + "primary_sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "sources": [ + "ops/mul_mat_q6_k_packed_f16_wmma.loom" + ], + "symbol": "ggml_mul_mat_q6_k_packed_token1_f16_wmma", + "workload_parameters": [ + { + "name": "token_count", + "type": "index" + } + ] + } + ], + "generated_count": 28, + "kernel_counts": { + "loom_libs:ggml_add_rmsnorm_binary_symmetric_i4_k32": 127, + "loom_libs:ggml_binary_f32": 1, + "loom_libs:ggml_binary_swiglu_symmetric_i4_k32": 64, + "loom_libs:ggml_flash_attention_f32_f16_wmma": 16, + "loom_libs:ggml_get_rows_f32": 98, + "loom_libs:ggml_mul_mat_q6_k_i8_prepacked_f16_wmma": 16, + "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma": 1, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma": 112, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma": 112, + "loom_libs:ggml_mul_mat_symmetric_i4_lowrow_wmma": 144, + "loom_libs:ggml_quantize_f32_symmetric_i4_k32": 16, + "loom_libs:ggml_rmsnorm_binary_f32": 1, + "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32": 1, + "loom_libs:ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32": 48, + "loom_libs:ggml_rmsnorm_mul_rope_f32": 32, + "loom_libs:ggml_set_rows": 32, + "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot": 48, + "loom_libs:llm_gated_delta_net_projection_epilogue_f32": 48, + "loom_libs:llm_ssm_conv_dconv4_silu_rollback_f32": 48 + }, + "loom_source": "tools/benchmarks/loom/v2_mtp4_program4.work.loom", + "model": "v2_mtp4_program4", + "scenario": "tg_c4", + "schema": "ggml-hrx-model-loom-benchmarks-v2" +} 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 new file mode 100644 index 000000000000..9f5922a7a8b3 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/loom/v2_mtp4_program4.work.loom @@ -0,0 +1,403 @@ +// Generated by tools/benchmarks/generate-model-benchmarks.py for v2_mtp4_program4 scenarios: tg_c4. +// Regenerate from an HRX command program dump rather than editing by hand. + +target.decl @ggml_binary_f32_gfx11_wave64 +target.decl @ggml_binary_swiglu_i4_gfx11_wave32 +target.decl @ggml_flash_attention_gfx11_wave64 +target.decl @ggml_get_rows_f32_gfx11_wave64 +target.decl @ggml_rmsnorm_binary_gfx11_wave32 +target.decl @ggml_rmsnorm_gfx11_wave32 +target.decl @ggml_set_rows_gfx11_wave64 + +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_add_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %lhs: buffer, %rhs: buffer, %residual_output: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_binary_f32_gfx11_wave64) @ggml_binary_f32(%element_count: index) launch(%element_count: index, %lhs: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_binary_swiglu_i4_gfx11_wave32) @ggml_binary_swiglu_symmetric_i4_k32() launch(%lhs: buffer, %rhs: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_flash_attention_gfx11_wave64) @ggml_flash_attention_f32_f16_wmma(%query_token_count: index, %key_value_token_count: index) launch(%query_token_count: index, %key_value_token_count: index, %query: buffer, %key: buffer, %value: buffer, %mask: buffer, %gate: buffer, %output: buffer) +kernel.decl target(@ggml_get_rows_f32_gfx11_wave64) @ggml_get_rows_f32(%token_count: index, %row_count: index, %hidden_size: index) launch(%token_count: index, %row_count: index, %hidden_size: index, %token_ids: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_q6_k_packed_token1_f16_wmma(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma() launch(%first_weight: buffer, %second_weight: buffer, %first_output: buffer, %second_output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer, %partial: buffer, %completion_counters: buffer) +kernel.decl @ggml_mul_mat_symmetric_i4_lowrow_wmma() launch(%weight: buffer, %input: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl @ggml_quantize_f32_symmetric_i4_k32() launch(%input: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_f32(%token_count: index) launch(%token_count: index, %input: buffer, %rhs: buffer, %output: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_binary_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +kernel.decl target(@ggml_rmsnorm_binary_gfx11_wave32) @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32(%token_count: index) launch(%token_count: index, %input: buffer, %weight: buffer, %raw_gate: buffer, %output: buffer, %quantized_values: buffer, %scales: buffer, %sums: buffer) +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(%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_c4 + +check.case public @v2_mtp4_program4_tg_c4_000_ggml_get_rows_f32_case { + %token_count = check.literal value(4) : index + %row_count = check.literal value(248320) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<16xi8> + %weight = check.generate.fill value(0) : tensor<715161600xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<16xi8>, tensor<715161600xi8>, tensor<81920xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_000_ggml_get_rows_f32_case> @v2_mtp4_program4_tg_c4_000_ggml_get_rows_f32 + +check.case public @v2_mtp4_program4_tg_c4_001_ggml_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_001_ggml_rmsnorm_binary_symmetric_i4_k32_case> @v2_mtp4_program4_tg_c4_001_ggml_rmsnorm_binary_symmetric_i4_k32 + +check.case public @v2_mtp4_program4_tg_c4_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<27033600xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + %partial = check.generate.fill value(0) : tensor<163840xi8> + %completion_counters = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<27033600xi8>, tensor<81920xi8>, tensor<163840xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>, tensor<163840xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @v2_mtp4_program4_tg_c4_002_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @v2_mtp4_program4_tg_c4_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<168960xi8> + %second_weight = check.generate.fill value(0) : tensor<168960xi8> + %first_output = check.generate.fill value(0) : tensor<768xi8> + %second_output = check.generate.fill value(0) : tensor<768xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<168960xi8>, tensor<168960xi8>, tensor<768xi8>, tensor<768xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @v2_mtp4_program4_tg_c4_003_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @v2_mtp4_program4_tg_c4_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<81920xi8>, tensor<98304xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_004_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @v2_mtp4_program4_tg_c4_004_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @v2_mtp4_program4_tg_c4_005_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(30720) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x30720xf32> + %output = check.generate.fill value(1.0) : tensor<1x30720xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x30720xf32>, tensor<1x30720xf32>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_005_ggml_get_rows_f32_case> @v2_mtp4_program4_tg_c4_005_ggml_get_rows_f32 + +check.case public @v2_mtp4_program4_tg_c4_006_llm_ssm_conv_dconv4_silu_rollback_f32_case { + %state = check.generate.fill value(0) : tensor<122880xi8> + %x = check.generate.fill value(0) : tensor<163840xi8> + %filter = check.generate.fill value(0) : tensor<163840xi8> + %output = check.generate.fill value(0) : tensor<163840xi8> + %cache0 = check.generate.fill value(0) : tensor<122880xi8> + %cache1 = check.generate.fill value(0) : tensor<122880xi8> + %cache2 = check.generate.fill value(0) : tensor<122880xi8> + %cache3 = check.generate.fill value(0) : tensor<122880xi8> + %cache4 = check.generate.fill value(0) : tensor<122880xi8> + kernel.launch @llm_ssm_conv_dconv4_silu_rollback_f32[](%state, %x, %filter, %output, %cache0, %cache1, %cache2, %cache3, %cache4) : [](tensor<122880xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<163840xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>, tensor<122880xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_006_llm_ssm_conv_dconv4_silu_rollback_f32_case> @v2_mtp4_program4_tg_c4_006_llm_ssm_conv_dconv4_silu_rollback_f32 + +check.case public @v2_mtp4_program4_tg_c4_007_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(20) : index + %hidden_size = check.literal value(786432) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<20x786432xf32> + %output = check.generate.fill value(1.0) : tensor<1x786432xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<20x786432xf32>, tensor<1x786432xf32>) + check.return +} + +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[%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 +} + +check.benchmark<@v2_mtp4_program4_tg_c4_008_llm_gated_delta_net_projection_epilogue_f32_case> @v2_mtp4_program4_tg_c4_008_llm_gated_delta_net_projection_epilogue_f32 + +check.case public @v2_mtp4_program4_tg_c4_009_llm_gated_delta_net_f32_wmma_head128_snapshot_case { + %q = check.generate.fill value(0) : tensor<131072xi8> + %k = check.generate.fill value(0) : tensor<131072xi8> + %v = check.generate.fill value(0) : tensor<147456xi8> + %g = check.generate.fill value(0) : tensor<768xi8> + %beta = check.generate.fill value(0) : tensor<768xi8> + %state_in = check.generate.fill value(0) : tensor<3145728xi8> + %snapshot_cache = check.generate.fill value(0) : tensor<40894464xi8> + %dst = check.generate.fill value(0) : tensor<98304xi8> + %rms_scales = check.generate.fill value(0) : tensor<768xi8> + kernel.launch @llm_gated_delta_net_f32_wmma_head128_snapshot[](%q, %k, %v, %g, %beta, %state_in, %snapshot_cache, %dst, %rms_scales) : [](tensor<131072xi8>, tensor<131072xi8>, tensor<147456xi8>, tensor<768xi8>, tensor<768xi8>, tensor<3145728xi8>, tensor<40894464xi8>, tensor<98304xi8>, tensor<768xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_009_llm_gated_delta_net_f32_wmma_head128_snapshot_case> @v2_mtp4_program4_tg_c4_009_llm_gated_delta_net_f32_wmma_head128_snapshot + +check.case public @v2_mtp4_program4_tg_c4_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case { + %token_count = check.literal value(192) : index + %input = check.generate.fill value(0) : tensor<98304xi8> + %weight = check.generate.fill value(0) : tensor<512xi8> + %raw_gate = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32[%token_count](%token_count, %input, %weight, %raw_gate, %output, %quantized_values, %scales, %sums) : [index](index, tensor<98304xi8>, tensor<512xi8>, tensor<98304xi8>, tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32_case> @v2_mtp4_program4_tg_c4_010_ggml_rmsnorm_gate_silu_mul_symmetric_i4_k32 + +check.case public @v2_mtp4_program4_tg_c4_011_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<16220160xi8> + %input = check.generate.fill value(0) : tensor<98304xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<16220160xi8>, tensor<98304xi8>, tensor<81920xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_011_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @v2_mtp4_program4_tg_c4_011_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @v2_mtp4_program4_tg_c4_012_ggml_add_rmsnorm_binary_symmetric_i4_k32_case { + %token_count = check.literal value(4) : index + %lhs = check.generate.fill value(0) : tensor<81920xi8> + %rhs = check.generate.fill value(0) : tensor<81920xi8> + %residual_output = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<20480xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_add_rmsnorm_binary_symmetric_i4_k32[%token_count](%token_count, %lhs, %rhs, %residual_output, %weight, %output, %quantized_values, %scales, %sums) : [index](index, tensor<81920xi8>, tensor<81920xi8>, tensor<81920xi8>, tensor<20480xi8>, tensor<81920xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_012_ggml_add_rmsnorm_binary_symmetric_i4_k32_case> @v2_mtp4_program4_tg_c4_012_ggml_add_rmsnorm_binary_symmetric_i4_k32 + +check.case public @v2_mtp4_program4_tg_c4_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case { + %first_weight = check.generate.fill value(0) : tensor<45957120xi8> + %second_weight = check.generate.fill value(0) : tensor<45957120xi8> + %first_output = check.generate.fill value(0) : tensor<278528xi8> + %second_output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma[](%first_weight, %second_weight, %first_output, %second_output, %quantized_values, %scales, %sums) : [](tensor<45957120xi8>, tensor<45957120xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma_case> @v2_mtp4_program4_tg_c4_013_ggml_mul_mat_symmetric_i4_lowrow_adjacent_dual_wmma + +check.case public @v2_mtp4_program4_tg_c4_014_ggml_binary_swiglu_symmetric_i4_k32_case { + %lhs = check.generate.fill value(0) : tensor<278528xi8> + %rhs = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<278528xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + kernel.launch @ggml_binary_swiglu_symmetric_i4_k32[](%lhs, %rhs, %output, %quantized_values, %scales, %sums) : [](tensor<278528xi8>, tensor<278528xi8>, tensor<278528xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_014_ggml_binary_swiglu_symmetric_i4_k32_case> @v2_mtp4_program4_tg_c4_014_ggml_binary_swiglu_symmetric_i4_k32 + +check.case public @v2_mtp4_program4_tg_c4_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case { + %weight = check.generate.fill value(0) : tensor<45957120xi8> + %input = check.generate.fill value(0) : tensor<278528xi8> + %output = check.generate.fill value(0) : tensor<81920xi8> + %quantized_values = check.generate.fill value(0) : tensor<34816xi8> + %scales = check.generate.fill value(0) : tensor<8704xi8> + %sums = check.generate.fill value(0) : tensor<8704xi8> + %partial = check.generate.fill value(0) : tensor<81920xi8> + %completion_counters = check.generate.fill value(0) : tensor<1280xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums, %partial, %completion_counters) : [](tensor<45957120xi8>, tensor<278528xi8>, tensor<81920xi8>, tensor<34816xi8>, tensor<8704xi8>, tensor<8704xi8>, tensor<81920xi8>, tensor<1280xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma_case> @v2_mtp4_program4_tg_c4_015_ggml_mul_mat_symmetric_i4_lowrow_split_k2_wmma + +check.case public @v2_mtp4_program4_tg_c4_016_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<32440320xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<196608xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<32440320xi8>, tensor<81920xi8>, tensor<196608xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_016_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @v2_mtp4_program4_tg_c4_016_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @v2_mtp4_program4_tg_c4_017_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(0) : tensor<81920xi8> + %weight = check.generate.fill value(0) : tensor<5611520xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_mul_mat_q6_k_i8_prepacked_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<81920xi8>, tensor<5611520xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_017_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma_case> @v2_mtp4_program4_tg_c4_017_ggml_mul_mat_q6_k_i8_prepacked_f16_wmma + +check.case public @v2_mtp4_program4_tg_c4_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case { + %weight = check.generate.fill value(0) : tensor<2703360xi8> + %input = check.generate.fill value(0) : tensor<81920xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + %quantized_values = check.generate.fill value(0) : tensor<10240xi8> + %scales = check.generate.fill value(0) : tensor<2560xi8> + %sums = check.generate.fill value(0) : tensor<2560xi8> + kernel.launch @ggml_mul_mat_symmetric_i4_lowrow_wmma[](%weight, %input, %output, %quantized_values, %scales, %sums) : [](tensor<2703360xi8>, tensor<81920xi8>, tensor<16384xi8>, tensor<10240xi8>, tensor<2560xi8>, tensor<2560xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_018_ggml_mul_mat_symmetric_i4_lowrow_wmma_case> @v2_mtp4_program4_tg_c4_018_ggml_mul_mat_symmetric_i4_lowrow_wmma + +check.case public @v2_mtp4_program4_tg_c4_019_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<195584xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<98304xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<195584xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<98304xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_019_ggml_rmsnorm_mul_rope_f32_case> @v2_mtp4_program4_tg_c4_019_ggml_rmsnorm_mul_rope_f32 + +check.case public @v2_mtp4_program4_tg_c4_020_ggml_rmsnorm_mul_rope_f32_case { + %input = check.generate.fill value(0) : tensor<16384xi8> + %weight = check.generate.fill value(0) : tensor<1024xi8> + %positions = check.generate.fill value(0) : tensor<64xi8> + %output = check.generate.fill value(0) : tensor<16384xi8> + kernel.launch @ggml_rmsnorm_mul_rope_f32[](%input, %weight, %positions, %output) : [](tensor<16384xi8>, tensor<1024xi8>, tensor<64xi8>, tensor<16384xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_020_ggml_rmsnorm_mul_rope_f32_case> @v2_mtp4_program4_tg_c4_020_ggml_rmsnorm_mul_rope_f32 + +check.case public @v2_mtp4_program4_tg_c4_021_ggml_set_rows_case { + %token_count = check.literal value(4) : index + %cache_row_count = check.literal value(262144) : index + %hidden_size = check.literal value(1024) : index + %rows = check.generate.fill value(0.0) : tensor<4x1024xf32> + %indices = check.generate.iota offset(0) step(1) period(262144) : tensor<4xi64> + %cache = check.generate.fill value(0.0) : tensor<262144x1024xf16> + kernel.launch @ggml_set_rows[%token_count, %cache_row_count, %hidden_size](%token_count, %cache_row_count, %hidden_size, %rows, %indices, %cache) : [index, index, index](index, index, index, tensor<4x1024xf32>, tensor<4xi64>, tensor<262144x1024xf16>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_021_ggml_set_rows_case> @v2_mtp4_program4_tg_c4_021_ggml_set_rows + +check.case public @v2_mtp4_program4_tg_c4_022_ggml_flash_attention_f32_f16_wmma_case { + %query_token_count = check.literal value(4) : index + %key_value_token_count = check.literal value(256) : index + %query = check.generate.fill value(0.0) : tensor<4x24x256xf32> + %key = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %value = check.generate.fill value(0.0) : tensor<256x4x256xf16> + %mask = check.generate.fill value(0.0) : tensor<4x256xf16> + %output = check.generate.fill value(1.0) : tensor<4x24x256xf32> + kernel.launch @ggml_flash_attention_f32_f16_wmma[%query_token_count, %key_value_token_count](%query_token_count, %key_value_token_count, %query, %key, %value, %mask, %output) : [index, index](index, index, tensor<4x24x256xf32>, tensor<256x4x256xf16>, tensor<256x4x256xf16>, tensor<4x256xf16>, tensor<4x24x256xf32>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_022_ggml_flash_attention_f32_f16_wmma_case> @v2_mtp4_program4_tg_c4_022_ggml_flash_attention_f32_f16_wmma + +check.case public @v2_mtp4_program4_tg_c4_023_ggml_quantize_f32_symmetric_i4_k32_case { + %input = check.generate.fill value(0) : tensor<98304xi8> + %quantized_values = check.generate.fill value(0) : tensor<12288xi8> + %scales = check.generate.fill value(0) : tensor<3072xi8> + %sums = check.generate.fill value(0) : tensor<3072xi8> + kernel.launch @ggml_quantize_f32_symmetric_i4_k32[](%input, %quantized_values, %scales, %sums) : [](tensor<98304xi8>, tensor<12288xi8>, tensor<3072xi8>, tensor<3072xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_023_ggml_quantize_f32_symmetric_i4_k32_case> @v2_mtp4_program4_tg_c4_023_ggml_quantize_f32_symmetric_i4_k32 + +check.case public @v2_mtp4_program4_tg_c4_024_ggml_binary_f32_case { + %element_count = check.literal value(20480) : index + %lhs = check.generate.fill value(2.0) : tensor<20480xf32> + %rhs = check.generate.fill value(3.0) : tensor<20480xf32> + %output = check.generate.fill value(0.0) : tensor<20480xf32> + kernel.launch @ggml_binary_f32[%element_count](%element_count, %lhs, %rhs, %output) : [index](index, tensor<20480xf32>, tensor<20480xf32>, tensor<20480xf32>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_024_ggml_binary_f32_case> @v2_mtp4_program4_tg_c4_024_ggml_binary_f32 + +check.case public @v2_mtp4_program4_tg_c4_025_ggml_rmsnorm_binary_f32_case { + %token_count = check.literal value(4) : index + %input = check.generate.fill value(2.0) : tensor<4x5120xf32> + %rhs = check.generate.fill value(3.0) : tensor<5120xf32> + %output = check.generate.fill value(0.0) : tensor<4x5120xf32> + kernel.launch @ggml_rmsnorm_binary_f32[%token_count](%token_count, %input, %rhs, %output) : [index](index, tensor<4x5120xf32>, tensor<5120xf32>, tensor<4x5120xf32>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_025_ggml_rmsnorm_binary_f32_case> @v2_mtp4_program4_tg_c4_025_ggml_rmsnorm_binary_f32 + +check.case public @v2_mtp4_program4_tg_c4_026_ggml_get_rows_f32_case { + %token_count = check.literal value(1) : index + %row_count = check.literal value(4) : index + %hidden_size = check.literal value(5120) : index + %token_ids = check.generate.fill value(0) : tensor<1xi32> + %weight = check.generate.fill value(0.0) : tensor<4x5120xf32> + %output = check.generate.fill value(1.0) : tensor<1x5120xf32> + kernel.launch @ggml_get_rows_f32[%token_count, %row_count, %hidden_size](%token_count, %row_count, %hidden_size, %token_ids, %weight, %output) : [index, index, index](index, index, index, tensor<1xi32>, tensor<4x5120xf32>, tensor<1x5120xf32>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_026_ggml_get_rows_f32_case> @v2_mtp4_program4_tg_c4_026_ggml_get_rows_f32 + +check.case public @v2_mtp4_program4_tg_c4_027_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case { + %token_count = check.literal value(1) : index + %input = check.generate.fill value(0) : tensor<20480xi8> + %weight = check.generate.fill value(0) : tensor<1042944000xi8> + %output = check.generate.fill value(0) : tensor<993280xi8> + kernel.launch @ggml_mul_mat_q6_k_packed_token1_f16_wmma[%token_count](%token_count, %input, %weight, %output) : [index](index, tensor<20480xi8>, tensor<1042944000xi8>, tensor<993280xi8>) + check.return +} + +check.benchmark<@v2_mtp4_program4_tg_c4_027_ggml_mul_mat_q6_k_packed_token1_f16_wmma_case> @v2_mtp4_program4_tg_c4_027_ggml_mul_mat_q6_k_packed_token1_f16_wmma diff --git a/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py new file mode 100755 index 000000000000..551fa6fe8fb9 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.py @@ -0,0 +1,531 @@ +#!/usr/bin/env python3 +# +# Run generated HRX model-scoped Loom benchmarks. + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import signal +import subprocess +from pathlib import Path +from typing import Any, NamedTuple + + +SCRIPT_DIR = Path(__file__).resolve().parent +HRX_DIR = SCRIPT_DIR.parents[1] +REPO_DIR = HRX_DIR.parents[2] +BENCHMARK_DIR = HRX_DIR / "benchmarks" + + +def fail(message: str) -> None: + raise SystemExit(message) + + +def load_json(path: Path) -> Any: + with path.open("r", encoding="utf-8") as f: + return json.load(f) + + +def find_tool(env_name: str, binary_name: str, build_dir: Path | None, build_suffix: str) -> Path: + override = os.environ.get(env_name) + if override: + path = Path(override) + if path.is_file() and os.access(path, os.X_OK): + return path + fail(f"{env_name} is not executable: {path}") + if build_dir is not None: + matches = sorted(build_dir.glob(build_suffix)) + for match in matches: + if match.is_file() and os.access(match, os.X_OK): + return match + found = shutil.which(binary_name) + if found: + return Path(found) + fail(f"could not find {binary_name}; set {env_name} or pass --build-dir") + + +def source_path(entry: dict[str, Any], source: str) -> Path: + corpus_dir = entry.get("corpus_dir") + if not corpus_dir: + fail(f"{entry['kernel']} has no corpus_dir in manifest") + 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", []) + return manifest.get("benchmarks", []) + + +def entry_status(entry: dict[str, Any]) -> str: + status = entry.get("status") + if status is not None: + return str(status) + if entry.get("benchmark"): + return "generated" + return "unsupported" + + +def entry_count(entry: dict[str, Any]) -> int: + return int(entry.get("count", entry.get("shape_multiplicity", 1))) + + +def parse_bool(value: str) -> bool: + value = value.lower() + if value in ("1", "true", "yes", "on"): + return True + if value in ("0", "false", "no", "off"): + return False + raise argparse.ArgumentTypeError(f"expected true or false, got {value!r}") + + +class CommandResult(NamedTuple): + state: str + error: str | None + returncode: int | None + signal_name: str | None + stdout: Path + stderr: Path + command: Path + + +def format_signal(returncode: int | None) -> str | None: + if returncode is None or returncode >= 0: + return None + signum = -returncode + try: + return signal.Signals(signum).name + except ValueError: + return f"SIG{signum}" + + +def process_output_text(output: str | bytes | None) -> str: + if output is None: + return "" + if isinstance(output, bytes): + return output.decode("utf-8", errors="replace") + return output + + +def run_command(command: list[str], cwd: Path, timeout_sec: float | None, bench_dir: Path, stage: str) -> CommandResult: + stdout_path = bench_dir / f"{stage}.stdout.txt" + stderr_path = bench_dir / f"{stage}.stderr.txt" + command_path = bench_dir / f"{stage}.command.json" + command_record: dict[str, Any] = { + "argv": command, + "cwd": str(cwd), + "timeout_sec": timeout_sec, + } + try: + result = subprocess.run(command, cwd=str(cwd), capture_output=True, text=True, timeout=timeout_sec) + except subprocess.TimeoutExpired as exc: + stdout_path.write_text(process_output_text(exc.stdout), encoding="utf-8") + stderr_path.write_text(process_output_text(exc.stderr), encoding="utf-8") + command_record.update( + { + "state": "timeout", + "returncode": None, + "signal": None, + "timed_out": True, + } + ) + command_path.write_text(json.dumps(command_record, indent=2, sort_keys=True) + "\n", encoding="utf-8") + return CommandResult( + "timeout", + f"{stage} timed out after {exc.timeout} seconds", + None, + None, + stdout_path, + stderr_path, + command_path, + ) + + stdout_path.write_text(result.stdout, encoding="utf-8") + stderr_path.write_text(result.stderr, encoding="utf-8") + sig_name = format_signal(result.returncode) + command_record.update( + { + "state": "ok" if result.returncode == 0 else "failed", + "returncode": result.returncode, + "signal": sig_name, + "timed_out": False, + } + ) + command_path.write_text(json.dumps(command_record, indent=2, sort_keys=True) + "\n", encoding="utf-8") + if result.returncode == 0: + return CommandResult("ok", None, result.returncode, sig_name, stdout_path, stderr_path, command_path) + if sig_name: + return CommandResult( + "failed", + f"{stage} exited with status {result.returncode} ({sig_name})", + result.returncode, + sig_name, + stdout_path, + stderr_path, + command_path, + ) + return CommandResult( + "failed", + f"{stage} exited with status {result.returncode}", + result.returncode, + sig_name, + stdout_path, + stderr_path, + command_path, + ) + + +def command_result_json(result: CommandResult) -> dict[str, Any]: + return { + "state": result.state, + "error": result.error, + "returncode": result.returncode, + "signal": result.signal_name, + "stdout": str(result.stdout), + "stderr": str(result.stderr), + "command": str(result.command), + } + + +def find_matching_char(text: str, open_index: int, open_char: str, close_char: str) -> int: + depth = 0 + for index in range(open_index, len(text)): + char = text[index] + if char == open_char: + depth += 1 + elif char == close_char: + depth -= 1 + if depth == 0: + return index + raise ValueError(f"unmatched {open_char!r}") + + +def parse_root_index_parameters(source_text: str, symbol: str) -> list[str]: + marker = f"@{symbol}(" + start = source_text.find(marker) + if start < 0: + raise ValueError(f"could not find root @{symbol}") + params_start = start + len(marker) - 1 + params_end = find_matching_char(source_text, params_start, "(", ")") + params_text = source_text[params_start + 1 : params_end] + names = [] + for raw_param in params_text.split(","): + parts = raw_param.strip().split() + if len(parts) < 2 or parts[-1] != "index": + continue + name = parts[-2].rstrip(":") + if not name.startswith("%"): + continue + names.append(name[1:]) + return names + + +def replace_index_uses(region: str, name: str, constant_name: str) -> str: + token = "%" + name + output = [] + index = 0 + while True: + match = region.find(token, index) + if match < 0: + output.append(region[index:]) + return "".join(output) + before_ok = match == 0 or not (region[match - 1].isalnum() or region[match - 1] == "_") + end = match + len(token) + after_ok = end == len(region) or not (region[end].isalnum() or region[end] == "_") + output.append(region[index:match]) + if before_ok and after_ok: + output.append("%" + constant_name) + else: + output.append(token) + index = end + + +def specialize_region(region: str, workload_values: list[tuple[str, int]]) -> str: + body = region + constants = [] + for name, value in workload_values: + constant_name = f"ggml_hrx_specialized_{name}" + constants.append(f" %{constant_name} = index.constant {value} : index\n") + body = replace_index_uses(body, name, constant_name) + return "\n" + "".join(constants) + body + + +def specialize_workload_text(source_text: str, symbol: str, workload_values: list[tuple[str, int]]) -> str: + if not workload_values: + return source_text + marker = f"@{symbol}(" + root_start = source_text.find(marker) + if root_start < 0: + raise ValueError(f"could not find root @{symbol}") + params_start = root_start + len(marker) - 1 + params_end = find_matching_char(source_text, params_start, "(", ")") + + config_open = source_text.find("{", params_end) + if config_open < 0: + raise ValueError("could not find kernel config block") + config_close = find_matching_char(source_text, config_open, "{", "}") + + launch_marker = "launch(" + launch_start = source_text.find(launch_marker, config_close) + if launch_start < 0: + raise ValueError("could not find launch region") + launch_body_open = source_text.find("{", launch_start) + if launch_body_open < 0: + raise ValueError("could not find launch body") + launch_body_close = find_matching_char(source_text, launch_body_open, "{", "}") + + config_body = source_text[config_open + 1 : config_close] + launch_body = source_text[launch_body_open + 1 : launch_body_close] + specialized_config = specialize_region(config_body, workload_values) + specialized_launch = specialize_region(launch_body, workload_values) + + return ( + source_text[: config_open + 1] + + specialized_config + + source_text[config_close: launch_body_open + 1] + + specialized_launch + + source_text[launch_body_close:] + ) + + +def specialize_linked_source(entry: dict[str, Any], linked_source: Path, output_source: Path) -> bool: + symbol = entry.get("symbol") or entry["kernel"].split(":")[-1] + integer_parameters = entry.get("integer_parameters", {}) + workload_parameters = entry.get("workload_parameters") or [] + workload_names = [param["name"] for param in workload_parameters if param.get("type") == "index"] + source_text = linked_source.read_text(encoding="utf-8") + if not workload_names: + workload_names = parse_root_index_parameters(source_text, symbol) + workload_values = [] + for name in workload_names: + if name not in integer_parameters: + continue + workload_values.append((name, int(integer_parameters[name]))) + if not workload_values: + return False + output_source.write_text(specialize_workload_text(source_text, symbol, workload_values), encoding="utf-8") + return True + + +def select_entries(manifest: dict[str, Any], benchmark: str | None, include_unsupported: bool) -> list[dict[str, Any]]: + entries = manifest_entries(manifest) + if benchmark: + requested = benchmark if benchmark.startswith("@") else "@" + benchmark + entries = [entry for entry in entries if entry.get("benchmark") == requested] + if not entries: + fail(f"benchmark not found in manifest: {requested}") + if include_unsupported: + return entries + return [entry for entry in entries if entry_status(entry) == "generated"] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Model slug, for example llama32_3b_f16.") + parser.add_argument("--scenario", required=True, help="Scenario slug, for example pp512.") + parser.add_argument("--build-dir", type=Path, help="llama.cpp build directory.") + parser.add_argument("--output-dir", type=Path, help="Directory for linked sources, plans, and results.") + parser.add_argument("--manifest", type=Path, help="Generated model scenario manifest.") + parser.add_argument("--benchmark", help="Single benchmark symbol to run.") + parser.add_argument("--device", default=os.environ.get("DEVICE", "amdgpu"), help="HAL device. Defaults to amdgpu.") + parser.add_argument("--dry-run-only", action="store_true", help="Stop after benchmark planning.") + parser.add_argument("--list", action="store_true", help="List selected generated benchmarks and exit.") + parser.add_argument("--include-unsupported", action="store_true", help="Include unsupported manifest entries when listing.") + parser.add_argument("--batch-size", default="64") + parser.add_argument("--iterations", default="10") + parser.add_argument("--warmup-iterations", default="1") + parser.add_argument("--continue-on-failure", action="store_true", help="Record failed benchmarks and continue running the rest.") + parser.add_argument("--benchmark-timeout-sec", type=float, default=300.0, help="Timeout per benchmark tool invocation. Defaults to 300.") + parser.add_argument("--profile-final-batch", type=parse_bool, default=True, help="Pass --profile-final-batch to iree-benchmark-loom. Defaults to true.") + parser.add_argument("--input-ring-count", help="Optional --input-ring-count value for iree-benchmark-loom.") + parser.add_argument("--benchmark-extra-arg", action="append", default=[], help="Additional argument to pass to benchmark invocations. May be repeated.") + parser.add_argument("--runtime-specialization", action=argparse.BooleanOptionalAction, default=True, help="Specialize workload index parameters after loom-link. Defaults to enabled.") + args = parser.parse_args() + + build_dir = args.build_dir + if build_dir is not None and not build_dir.is_absolute(): + build_dir = REPO_DIR / build_dir + manifest_path = args.manifest or (BENCHMARK_DIR / "loom" / f"{args.model}.{args.scenario}.json") + manifest = load_json(manifest_path) + loom_source = HRX_DIR / manifest["loom_source"] + entries = select_entries(manifest, args.benchmark, args.include_unsupported) + + if args.list: + for entry in entries: + count = entry_count(entry) + print(f"{entry.get('benchmark') or '-'} {entry_status(entry)} {entry['kernel']} x{count}") + return + + entries = [entry for entry in entries if entry_status(entry) == "generated"] + if not entries: + fail("no generated benchmarks selected") + if args.output_dir is None: + fail("--output-dir is required unless --list is used") + + loom_link = find_tool("LOOM_LINK", "loom-link", build_dir, "**/loom-link") + iree_benchmark_loom = find_tool("IREE_BENCHMARK_LOOM", "iree-benchmark-loom", build_dir, "**/iree-benchmark-loom") + args.output_dir.mkdir(parents=True, exist_ok=True) + summary_path = args.output_dir / "results.jsonl" + with summary_path.open("w", encoding="utf-8") as summary: + for entry in entries: + benchmark = entry["benchmark"] + safe_name = benchmark[1:] + bench_dir = args.output_dir / safe_name + artifact_dir = bench_dir / "artifacts" + bench_dir.mkdir(parents=True, exist_ok=True) + artifact_dir.mkdir(parents=True, exist_ok=True) + linked_source = bench_dir / "linked.loom" + specialized_source = bench_dir / "linked.runtime-specialized.loom" + plan_output = bench_dir / "plan.json" + result_output = bench_dir / "results.json" + stage_results: dict[str, Any] = {} + config_args = [f"--config={name}={value}" for name, value in sorted(entry.get("compile_parameters", {}).items())] + + link_inputs = loom_link_source_args(entry) + state = "ok" + error = None + link_result = run_command( + [ + str(loom_link), + str(loom_source), + *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", + ) + 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": + benchmark_source = linked_source + runtime_specialized = False + if args.runtime_specialization: + try: + runtime_specialized = specialize_linked_source(entry, linked_source, specialized_source) + except ValueError as exc: + state = "failed" + error = f"runtime specialization failed: {exc}" + if runtime_specialized: + benchmark_source = specialized_source + else: + runtime_specialized = False + else: + benchmark_source = linked_source + runtime_specialized = False + if state == "ok": + dry_run = [ + str(iree_benchmark_loom), + str(benchmark_source), + f"--benchmark={benchmark}", + f"--device={args.device}", + "--dry-run", + "--measure=dispatch_complete", + "--output-format=jsonl", + f"--output={plan_output}", + *config_args, + ] + dry_run_result = run_command(dry_run, REPO_DIR, args.benchmark_timeout_sec, bench_dir, "dry-run") + stage_results["dry-run"] = command_result_json(dry_run_result) + state = dry_run_result.state + error = dry_run_result.error + if state == "ok" and not args.dry_run_only: + benchmark_command = [ + str(iree_benchmark_loom), + str(benchmark_source), + f"--benchmark={benchmark}", + f"--device={args.device}", + "--measure=dispatch_complete", + f"--batch-size={args.batch_size}", + f"--iterations={args.iterations}", + f"--warmup-iterations={args.warmup_iterations}", + f"--profile-final-batch={str(args.profile_final_batch).lower()}", + f"--artifact-bundle-dir={artifact_dir}", + "--output-format=jsonl", + f"--output={result_output}", + *config_args, + *args.benchmark_extra_arg, + ] + if args.input_ring_count is not None: + benchmark_command.append(f"--input-ring-count={args.input_ring_count}") + benchmark_result = run_command(benchmark_command, REPO_DIR, args.benchmark_timeout_sec, bench_dir, "benchmark") + stage_results["benchmark"] = command_result_json(benchmark_result) + state = benchmark_result.state + error = benchmark_result.error + summary.write( + json.dumps( + { + "benchmark": benchmark, + "kernel": entry["kernel"], + "count": entry_count(entry), + "linked_source": str(linked_source), + "benchmark_source": str(benchmark_source), + "runtime_specialization": runtime_specialized, + "plan": str(plan_output), + "results": str(result_output) if not args.dry_run_only else None, + "state": state, + "error": error, + "stages": stage_results, + }, + sort_keys=True, + ) + + "\n" + ) + summary.flush() + print(f"{state} {benchmark}") + if state != "ok" and not args.continue_on_failure: + fail(f"{benchmark}: {error}") + print(f"wrote {summary_path}") + + +if __name__ == "__main__": + main() diff --git a/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.sh b/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.sh new file mode 100755 index 000000000000..b5510cbdae83 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/run-model-benchmarks.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +set -euo pipefail + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +exec python3 "${script_dir}/run-model-benchmarks.py" "$@" diff --git a/ggml/src/ggml-hrx/tools/benchmarks/summarize-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/summarize-model-benchmarks.py new file mode 100755 index 000000000000..b1ec1255977b --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/summarize-model-benchmarks.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +# +# Summarize generated HRX model-scoped Loom benchmark results. + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def fail(message: str) -> None: + raise SystemExit(message) + + +def load_jsonl(path: Path) -> list[dict[str, Any]]: + rows = [] + with path.open("r", encoding="utf-8") as f: + for line in f: + if line.strip(): + rows.append(json.loads(line)) + return rows + + +def load_benchmark_result(path: Path) -> dict[str, Any] | None: + with path.open("r", encoding="utf-8", errors="replace") as f: + for line in f: + if not line.strip(): + continue + row = json.loads(line) + if row.get("row") == "benchmark": + return row.get("benchmark_result") + return None + + +def timing_from_result(result: dict[str, Any], metric: str) -> float | None: + if metric == "operation-p50": + return result.get("measurement", {}).get("operation_timing_ns", {}).get("p50") + if metric == "profile-p50": + return ( + result.get("profile_replay", {}) + .get("dispatch_timing", {}) + .get("dispatch_distribution", {}) + .get("duration_ns", {}) + .get("p50") + ) + fail(f"unsupported metric: {metric}") + + +def secondary_profile_p50(result: dict[str, Any]) -> float | None: + return ( + result.get("profile_replay", {}) + .get("dispatch_timing", {}) + .get("dispatch_distribution", {}) + .get("duration_ns", {}) + .get("p50") + ) + + +def summarize(rows: list[dict[str, Any]], metric: str) -> dict[str, Any]: + shape_rows = [] + for row in rows: + count = int(row.get("count", row.get("shape_multiplicity", 1))) + result_path = row.get("results") + shape = { + "benchmark": row.get("benchmark"), + "kernel": row.get("kernel"), + "count": count, + "state": row.get("state", "ok"), + "error": row.get("error"), + "metric_ns": None, + "profile_p50_ns": None, + "weighted_ns": None, + } + if shape["state"] != "ok": + shape_rows.append(shape) + continue + if not result_path: + shape["state"] = "missing_results" + shape["error"] = "runner did not record a results path" + shape_rows.append(shape) + continue + result_file = Path(result_path) + if not result_file.is_file(): + shape["state"] = "missing_results" + shape["error"] = f"results file not found: {result_file}" + shape_rows.append(shape) + continue + result = load_benchmark_result(result_file) + if result is None: + shape["state"] = "missing_benchmark_row" + shape["error"] = "results file did not contain a benchmark row" + shape_rows.append(shape) + continue + if result.get("state") != "ok": + shape["state"] = result.get("state", "failed") + shape["error"] = result.get("failure", {}).get("kind") or "benchmark did not complete successfully" + shape_rows.append(shape) + continue + metric_ns = timing_from_result(result, metric) + if metric_ns is None: + shape["state"] = "missing_metric" + shape["error"] = f"metric not found: {metric}" + shape_rows.append(shape) + continue + shape["metric_ns"] = metric_ns + shape["profile_p50_ns"] = secondary_profile_p50(result) + shape["weighted_ns"] = metric_ns * count + shape_rows.append(shape) + + kernels: dict[str, dict[str, Any]] = {} + for shape in shape_rows: + kernel = shape["kernel"] + entry = kernels.setdefault( + kernel, + { + "kernel": kernel, + "shape_count": 0, + "invocation_count": 0, + "measured_shape_count": 0, + "measured_invocation_count": 0, + "missing_or_failed_count": 0, + "missing_or_failed_invocations": 0, + "weighted_ns": 0.0, + "weighted_percent": 0.0, + "shapes": [], + }, + ) + entry["shape_count"] += 1 + entry["invocation_count"] += shape["count"] + entry["shapes"].append(shape) + if shape["weighted_ns"] is None: + entry["missing_or_failed_count"] += 1 + entry["missing_or_failed_invocations"] += shape["count"] + continue + entry["measured_shape_count"] += 1 + entry["measured_invocation_count"] += shape["count"] + entry["weighted_ns"] += shape["weighted_ns"] + + total_weighted_ns = sum(entry["weighted_ns"] for entry in kernels.values()) + for entry in kernels.values(): + if total_weighted_ns: + entry["weighted_percent"] = entry["weighted_ns"] / total_weighted_ns * 100.0 + kernel_rows = sorted(kernels.values(), key=lambda entry: entry["weighted_ns"], reverse=True) + return { + "schema": "ggml-hrx-model-loom-benchmark-summary-v2", + "metric": metric, + "metric_description": "operation_timing_ns.p50" if metric == "operation-p50" else "profile_replay dispatch_distribution duration_ns p50", + "shape_count": len(shape_rows), + "measured_shape_count": sum(1 for shape in shape_rows if shape["weighted_ns"] is not None), + "missing_or_failed_shape_count": sum(1 for shape in shape_rows if shape["weighted_ns"] is None), + "invocation_count": sum(shape["count"] for shape in shape_rows), + "measured_invocation_count": sum(shape["count"] for shape in shape_rows if shape["weighted_ns"] is not None), + "missing_or_failed_invocations": sum(shape["count"] for shape in shape_rows if shape["weighted_ns"] is None), + "total_weighted_ns": total_weighted_ns, + "kernels": kernel_rows, + } + + +def write_markdown(path: Path, summary: dict[str, Any]) -> None: + lines = [ + "# HRX Model Loom Benchmark Summary", + "", + f"Metric: `{summary['metric_description']}`", + "", + f"Measured shapes: {summary['measured_shape_count']} / {summary['shape_count']}", + f"Measured invocations: {summary['measured_invocation_count']} / {summary['invocation_count']}", + "", + "| Kernel | Shapes | Invocations | Measured Shapes | Weighted ms | Weighted Share | Missing/Failed Shapes |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for kernel in summary["kernels"]: + weighted_ms = kernel["weighted_ns"] / 1_000_000.0 + lines.append( + f"| `{kernel['kernel']}` | {kernel['shape_count']} | {kernel['invocation_count']} | " + f"{kernel['measured_shape_count']} | {weighted_ms:.3f} | " + f"{kernel['weighted_percent']:.2f}% | {kernel['missing_or_failed_count']} |" + ) + lines.append("") + lines.append("## Missing Or Failed Shapes") + lines.append("") + lines.append("| Kernel | Benchmark | Invocations | State | Error |") + lines.append("| --- | --- | ---: | --- | --- |") + missing = False + for kernel in summary["kernels"]: + for shape in kernel["shapes"]: + if shape["weighted_ns"] is not None: + continue + missing = True + error = str(shape.get("error") or "").replace("|", "\\|") + lines.append( + f"| `{shape['kernel']}` | `{shape['benchmark']}` | {shape['count']} | " + f"`{shape['state']}` | {error} |" + ) + if not missing: + lines.append("| - | - | 0 | - | - |") + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("results_jsonl", type=Path, help="Runner results.jsonl path.") + parser.add_argument("--metric", choices=["operation-p50", "profile-p50"], default="operation-p50") + parser.add_argument("--output-json", type=Path, help="Summary JSON output path.") + parser.add_argument("--output-md", type=Path, help="Summary Markdown output path.") + args = parser.parse_args() + + rows = load_jsonl(args.results_jsonl) + summary = summarize(rows, args.metric) + output_json = args.output_json or (args.results_jsonl.parent / "summary.json") + output_md = args.output_md or (args.results_jsonl.parent / "summary.md") + output_json.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") + write_markdown(output_md, summary) + print(f"wrote {output_json}") + print(f"wrote {output_md}") + + +if __name__ == "__main__": + main() diff --git a/ggml/src/ggml-hrx/tools/benchmarks/test-generate-model-benchmarks.py b/ggml/src/ggml-hrx/tools/benchmarks/test-generate-model-benchmarks.py new file mode 100644 index 000000000000..097adb476524 --- /dev/null +++ b/ggml/src/ggml-hrx/tools/benchmarks/test-generate-model-benchmarks.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import importlib.util +import sys +import unittest +from pathlib import Path + + +SCRIPT = Path(__file__).with_name("generate-model-benchmarks.py") +SPEC = importlib.util.spec_from_file_location("generate_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) + + +def gated_flash_attention_command(gate_bytes: int = 48128) -> dict: + return { + "ordinal": 59, + "kernel": "loom_libs:ggml_flash_attention_f32_f16_wmma", + "integer_parameters": { + "query_token_count": 1, + "key_value_token_count": 256, + }, + "compile_parameters": { + "ggml.flash_attention.apply_gate": "1", + "ggml.flash_attention.query_head_count": "24", + "ggml.flash_attention.key_value_head_count": "4", + "ggml.flash_attention.qk_head_size": "256", + "ggml.flash_attention.value_head_size": "256", + "ggml.flash_attention.gate_stride_head": "512", + "ggml.flash_attention.gate_stride_token": "12288", + }, + "bindings": [ + {"name": "query", "length": 24576}, + {"name": "key", "length": 524288}, + {"name": "value", "length": 524288}, + {"name": "mask", "length": 512}, + {"name": "gate", "length": gate_bytes}, + {"name": "output", "length": 24576}, + ], + } + + +class FlashAttentionCaseTest(unittest.TestCase): + def test_materializes_distinct_gate_with_exact_binding_capacity(self) -> None: + case = MODULE.case_flash_attention( + "qwen38_tg8_flash_attention", + gated_flash_attention_command(), + "ggml_flash_attention_f32_f16_wmma", + "ggml", + ) + + self.assertIn("%gate = check.generate.fill value(0.0) : tensor<12032xf32>", case) + self.assertIn("%mask, %gate, %output)", case) + self.assertIn("tensor<12032xf32>", case) + self.assertNotIn("%mask, %query, %output)", case) + + def test_rejects_gate_binding_smaller_than_strided_access(self) -> None: + with self.assertRaisesRegex(SystemExit, "strided access requires 12032"): + MODULE.case_flash_attention( + "qwen38_tg8_flash_attention", + gated_flash_attention_command(gate_bytes=24576), + "ggml_flash_attention_f32_f16_wmma", + "ggml", + ) + + def test_rejects_gate_binding_not_aligned_to_f32(self) -> None: + with self.assertRaisesRegex(SystemExit, "non-f32 gate binding length"): + MODULE.case_flash_attention( + "qwen38_tg8_flash_attention", + gated_flash_attention_command(gate_bytes=48127), + "ggml_flash_attention_f32_f16_wmma", + "ggml", + ) + + +if __name__ == "__main__": + unittest.main() 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/ggml/src/ggml-hrx/tools/check_hrx_revision.py b/ggml/src/ggml-hrx/tools/check_hrx_revision.py new file mode 100644 index 000000000000..9ce6a441e2da --- /dev/null +++ b/ggml/src/ggml-hrx/tools/check_hrx_revision.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Check an HRX checkout against the backend's recorded minimum commit. + +CMake and callers using installed packages share this local ancestry check. +Success prints the accepted revision; failure prints context and exits nonzero. +No history is fetched and no source files are changed. +""" + +import argparse +from pathlib import Path +import subprocess +import sys + + +def check_revision(source: Path) -> int: + minimum = (Path(__file__).resolve().parents[1] / "hrx-minimum-commit.txt").read_text(encoding="utf-8").strip() + head = "unknown" + + def git(*args: str) -> str: + return subprocess.run( + ["git", "-C", str(source), *args], check=True, text=True, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, + ).stdout.strip() + + try: + head = git("rev-parse", "--verify", "HEAD^{commit}") + git("cat-file", "-e", f"{minimum}^{{commit}}") + git("merge-base", "--is-ancestor", minimum, head) + except (OSError, subprocess.CalledProcessError) as error: + if isinstance(error, subprocess.CalledProcessError): + reason = error.stderr.strip() or "HRX HEAD does not have the required commit in its available ancestry." + else: + reason = str(error) + print( + f"{reason}\nHRX_SOURCE_DIR: {source}\nRequired HRX commit: {minimum}\nActual HRX HEAD: {head}\n" + "Use an HRX Git checkout at the required commit or a descendant. Ensure Git is installed, " + "fetch missing history manually (unshallow a shallow clone if needed), then rerun the check.", + file=sys.stderr, + ) + return 1 + + print(f"HRX revision {head} satisfies minimum {minimum}") + return 0 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source", type=Path, help="HRX source checkout") + args = parser.parse_args() + return check_revision(args.source.resolve()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py b/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py index 31da06085261..e7104c635cb0 100755 --- a/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py +++ b/ggml/src/ggml-hrx/tools/compile_qwen_kernel_corpus.py @@ -19,7 +19,16 @@ ATTR_RE = re.compile(r"(?P[A-Za-z0-9_.]+)\s*=\s*(?P-?[0-9]+)") CASE_RE = re.compile(r"check\.case(?:\s+public)?\s+@(?P[A-Za-z0-9_]+)\s*\{") LITERAL_RE = re.compile(r"%(?P[A-Za-z0-9_]+)\s*=\s*check\.literal\s+value\((?P-?[0-9]+)\)\s*:\s*index") -FUNC_CALL_RE = re.compile(r"func\.call\s+@(?P[A-Za-z0-9_]+)\s*\(") +KERNEL_LAUNCH_RE = re.compile( + r"kernel\.launch\s+@(?P[A-Za-z0-9_]+)" + r"(?:\[(?P[^]]*)\])?\s*\(" +) +SSA_VALUE_RE = re.compile(r"%(?P[A-Za-z0-9_]+)") +TARGET_DEF_RE = re.compile(r"amdgpu\.target<(?P[^>]+)>\s+@(?P[A-Za-z0-9_]+)") +KERNEL_TARGET_RE = re.compile( + r"kernel\.def[^\n{]*target\(@(?P[A-Za-z0-9_]+)\)[^\n{]*@(?P[A-Za-z0-9_]+)\s*\(" +) +KERNEL_DEF_RE = re.compile(r"kernel\.def[^\n{]*@(?P[A-Za-z0-9_]+)\s*\(") def run(command: list[str], *, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: @@ -49,8 +58,9 @@ def main() -> int: manifest = json.loads(args.manifest.read_text(encoding="utf-8")) revision = require_run(["git", "-C", str(args.hrx_source), "rev-parse", "HEAD"]).stdout.strip() - if revision != manifest["upstream_revision"]: - raise RuntimeError(f"HRX checkout {revision} does not match corpus {manifest['upstream_revision']}") + upstream_revision = manifest.get("upstream_revision", "local") + if upstream_revision != "local" and revision != upstream_revision: + raise RuntimeError(f"HRX checkout {revision} does not match corpus {upstream_revision}") if require_run(["git", "-C", str(args.hrx_source), "status", "--porcelain"]).stdout.strip(): raise RuntimeError("refusing to compile against a dirty pinned HRX checkout") override = os.environ.get("HSA_OVERRIDE_GFX_VERSION") @@ -65,12 +75,34 @@ def main() -> int: linked_dir.mkdir(exist_ok=True) linked_paths: dict[str, pathlib.Path] = {} + def module_inputs(name: str) -> list[pathlib.Path]: + paths: list[pathlib.Path] = [] + seen: set[pathlib.Path] = set() + + def add(path: pathlib.Path) -> None: + if path not in seen: + seen.add(path) + paths.append(path) + + def visit(module_name: str) -> None: + recipe = modules[module_name] + for source in recipe["srcs"]: + add(args.corpus_dir / source) + for library in recipe["libraries"]: + if library.startswith(":"): + visit(library[1:]) + else: + add(args.corpus_dir / library) + + visit(name) + return paths + def link_module(name: str) -> pathlib.Path: if name in linked_paths: return linked_paths[name] recipe = modules[name] output = linked_dir / f"{name}.loom" - command = [str(args.loom_link), "--mode=archive", f"--output={output}"] + command = [str(args.loom_link), "--mode=merge", f"--output={output}"] command.extend(str(args.corpus_dir / source) for source in recipe["srcs"]) for library in recipe["libraries"]: path = link_module(library[1:]) if library.startswith(":") else args.corpus_dir / library @@ -80,20 +112,45 @@ def link_module(name: str) -> pathlib.Path: return output exports = {item["symbol"]: item for item in manifest["exports"]} + kernel_defs_by_source: dict[pathlib.Path, set[str]] = {} + + def kernel_defs(path: pathlib.Path) -> set[str]: + if path not in kernel_defs_by_source: + text = path.read_text(encoding="utf-8") + kernel_defs_by_source[path] = {match.group("name") for match in KERNEL_DEF_RE.finditer(text)} + return kernel_defs_by_source[path] - def source_for_root(root: str) -> pathlib.Path: - source_name = exports[root]["source"] - direct_modules = [name for name, recipe in modules.items() if source_name in recipe["srcs"]] - if direct_modules: - return link_module(direct_modules[0]) - library_modules = [name for name, recipe in modules.items() if source_name in recipe["libraries"]] - if library_modules: - return link_module(library_modules[0]) - return args.corpus_dir / source_name + def internal_inputs(root: str, inputs: list[pathlib.Path]) -> list[pathlib.Path]: + owners = {path for path in inputs if root in kernel_defs(path)} + if not owners: + return inputs + return [path for path in inputs if path in owners or not kernel_defs(path)] + + def export_inputs(root: str) -> list[pathlib.Path]: + recipe = exports[root]["compile_recipe"] + return [ + args.corpus_dir / source + for source in recipe["primary_sources"] + recipe["library_sources"] + ] + + def link_root(root: str, configs: tuple[str, ...] | list[str], inputs: list[pathlib.Path], output: pathlib.Path) -> None: + output.parent.mkdir(parents=True, exist_ok=True) + target = "amdgpu:gfx11-generic" if args.target == "gfx1100" else "amdgpu:gfx1151" + command = [ + str(args.loom_link), + "--mode=link", + f"--root=@{root}", + f"--target={target}", + "--require-resolved-config", + "--to=text", + f"--output={output}", + ] + command.extend(f"--config={config}" for config in configs) + command.extend(str(path) for path in inputs) + require_run(command) results: list[dict[str, object]] = [] compiled_keys: set[tuple[str, tuple[int, ...], tuple[str, ...]]] = set() - root_sources: dict[str, pathlib.Path] = {} planned_invocation_count = 0 resolved_invocation_count = 0 for case in manifest["plan_cases"]: @@ -114,13 +171,20 @@ def source_for_root(root: str) -> pathlib.Path: if not plan_rows: raise RuntimeError(f"BUILD recipe {case['name']} produced no planner rows") source_text = source.read_text(encoding="utf-8") + target_selectors = { + match.group("name"): match.group("selector") for match in TARGET_DEF_RE.finditer(source_text) + } + kernel_selectors = { + match.group("name"): target_selectors.get(match.group("target"), "") + for match in KERNEL_TARGET_RE.finditer(source_text) + } benchmark_defs = { match.group("name"): (match.group("case"), {item.group("name"): int(item.group("value")) for item in ATTR_RE.finditer(match.group("attrs") or "")}) for match in BENCHMARK_RE.finditer(source.read_text(encoding="utf-8")) } case_literals: dict[str, dict[str, int]] = {} - case_calls: dict[str, list[str]] = {} + case_invocations: dict[str, list[tuple[str, tuple[str, ...]]]] = {} for match in CASE_RE.finditer(source_text): depth = 1 cursor = match.end() @@ -132,43 +196,63 @@ def source_for_root(root: str) -> pathlib.Path: case_literals[match.group("name")] = { item.group("name"): int(item.group("value")) for item in LITERAL_RE.finditer(body) } - case_calls[match.group("name")] = [item.group("name") for item in FUNC_CALL_RE.finditer(body)] + case_invocations[match.group("name")] = [ + ( + item.group("name"), + tuple(value.group("name") for value in SSA_VALUE_RE.finditer(item.group("workload") or "")), + ) + for item in KERNEL_LAUNCH_RE.finditer(body) + ] configs = [item.removeprefix("--config=") for item in case["args"] if item.startswith("--config=")] selects_benchmark = any(item.startswith("--benchmark=") for item in case["args"]) owned_sources = set(modules[case["link_module"]]["srcs"]) if "link_module" in case else {case["source"]} + case_inputs = module_inputs(case["link_module"]) if "link_module" in case else [args.corpus_dir / case["source"]] for row in plan_rows: benchmark_case, attrs = benchmark_defs.get(row["benchmark"], (row["case"], {})) concrete_values = dict(case_literals.get(benchmark_case, {})) concrete_values.update(attrs) - expected_count = int(row.get("actual_invocation_count", 1)) + expected_count = int(row.get("kernel_launch_count", row.get("actual_invocation_count", 1))) planned_invocation_count += expected_count - if row.get("actual_entry"): - roots = [row["actual_entry"]] - else: - roots = [root for root in case_calls.get(benchmark_case, []) if root in exports] - if len(roots) != expected_count: + invocations = case_invocations.get(benchmark_case, []) + actual_entry = row.get("kernel_entry", row.get("actual_entry")) + if actual_entry: + invocations = [item for item in invocations if item[0] == actual_entry] + if not invocations and actual_entry in exports: + invocations = [ + ( + actual_entry, + tuple(item["name"] for item in exports[actual_entry]["workload_parameters"]), + ) + ] + if len(invocations) != expected_count: raise RuntimeError( f"BUILD recipe {case['name']} planner reports {expected_count} invocations for " - f"{benchmark_case}, but source resolves {len(roots)} exported calls: {roots}" + f"{benchmark_case}, but source resolves {len(invocations)} kernel launches: " + f"{[item[0] for item in invocations]}" ) - resolved_invocation_count += len(roots) - for root in roots: - if root not in exports: - raise RuntimeError(f"BUILD recipe {case['name']} selected unmanifested kernel {root}") - if not selects_benchmark and exports[root]["source"] not in owned_sources: + resolved_invocation_count += len(invocations) + for root, parameter_names in invocations: + if not selects_benchmark and root in exports and exports[root]["source"] not in owned_sources: continue - root_sources.setdefault(root, source) - parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] missing = [name for name in parameter_names if name not in concrete_values] if missing: raise RuntimeError(f"{case['name']} recipe omits concrete {missing} for {root}") workload = tuple(concrete_values[name] for name in parameter_names) - key = (root, workload, tuple(configs)) + key = (root, workload, tuple(sorted(configs))) if key in compiled_keys: continue compiled_keys.add(key) compile_dir = case_dir / f"{root}-{len(compiled_keys):03d}" - command = [str(args.compiler), "--target", args.target, "--source", str(source), + selector = kernel_selectors.get(root, "") + if selector and selector not in ("gfx11-generic", args.target): + results.append({"case": case["name"], "root": root, "workload": workload, + "configs": configs, "artifact_dir": str(compile_dir), + "status": "skipped-incompatible-target", "target_selector": selector}) + continue + linked_root = compile_dir / "linked.loom" + inputs = export_inputs(root) if root in exports else internal_inputs(root, case_inputs) + link_root(root, configs, inputs, linked_root) + command = [str(args.compiler), "--target", args.target, "--source", str(linked_root), "--root", root, "--output", str(compile_dir)] for config in configs: command.extend(("--config", config)) @@ -188,59 +272,83 @@ def source_for_root(root: str) -> pathlib.Path: schedule = json.loads(schedule_path.read_text(encoding="utf-8")) schedule_dir = args.output_dir / "schedules" / schedule_path.parent.name schedule_dir.mkdir(parents=True, exist_ok=True) - for invocation in schedule.get("invocations", []): - for dispatch in invocation.get("dispatches", []): - specialization = dispatch["kernel"] - if specialization.get("execution") != "native": - raise RuntimeError(f"schedule {schedule_path} contains non-native dispatch {specialization['variant']}") - root = specialization["variant"] - if root not in exports: - raise RuntimeError(f"schedule {schedule_path} selects unmanifested kernel {root}") - schedule_requirement_count += 1 - parameters = specialization.get("parameters", {}) - parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] - missing = [name for name in parameter_names if name not in parameters] - if missing: - raise RuntimeError(f"schedule {schedule_path} omits concrete {missing} for {root}") - workload = tuple(int(parameters[name]) for name in parameter_names) - configs = tuple(f"{name}={value}" for name, value in - sorted(specialization.get("compile_parameters", {}).items())) - key = (root, workload, configs) - if key in compiled_keys: + specializations: list[dict[str, object]] = [] + if schedule.get("schema") == "ggml-hrx-command-program-v1": + for command in schedule.get("commands", []): + if command.get("kind") != "Kernel": continue - schedule_unique_requirement_count += 1 - compiled_keys.add(key) - source = root_sources[root] if root in root_sources else source_for_root(root) - compile_dir = schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}" - command = [str(args.compiler), "--target", args.target, "--source", str(source), - "--root", root, "--output", str(compile_dir)] - for config in configs: - command.extend(("--config", config)) - for value in workload: - command.extend(("--workload", str(value))) - compile_result = run(command, env=os.environ.copy()) - (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stdout.txt").write_text( - compile_result.stdout, encoding="utf-8") - (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stderr.txt").write_text( - compile_result.stderr, encoding="utf-8") - results.append({"case": f"schedule:{schedule_path}", "root": root, "workload": workload, - "configs": configs, "artifact_dir": str(compile_dir), - "status": "ok" if compile_result.returncode == 0 else "failed"}) + qualified_root = str(command.get("kernel", "")) + if not qualified_root.startswith("qwen3_moe:"): + continue + specializations.append({ + "variant": qualified_root.removeprefix("qwen3_moe:"), + "parameters": command.get("integer_parameters", {}), + "compile_parameters": command.get("compile_parameters", {}), + }) + else: + for invocation in schedule.get("invocations", []): + for dispatch in invocation.get("dispatches", []): + specialization = dispatch["kernel"] + if specialization.get("execution") != "native": + raise RuntimeError( + f"schedule {schedule_path} contains non-native dispatch {specialization['variant']}" + ) + specializations.append(specialization) + if not specializations: + raise RuntimeError(f"schedule {schedule_path} contains no Qwen MoE kernel requirements") + for specialization in specializations: + root = str(specialization["variant"]) + if root not in exports: + raise RuntimeError(f"schedule {schedule_path} selects unmanifested Qwen MoE kernel {root}") + schedule_requirement_count += 1 + parameters = specialization.get("parameters", {}) + parameter_names = [item["name"] for item in exports[root]["workload_parameters"]] + missing = [name for name in parameter_names if name not in parameters] + if missing: + raise RuntimeError(f"schedule {schedule_path} omits concrete {missing} for {root}") + workload = tuple(int(parameters[name]) for name in parameter_names) + compile_parameters = specialization.get("compile_parameters", {}) + configs = tuple(f"{name}={value}" for name, value in sorted(compile_parameters.items())) + key = (root, workload, configs) + if key in compiled_keys: + continue + schedule_unique_requirement_count += 1 + compiled_keys.add(key) + compile_dir = schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}" + linked_root = compile_dir / "linked.loom" + link_root(root, configs, export_inputs(root), linked_root) + command = [str(args.compiler), "--target", args.target, "--source", str(linked_root), + "--root", root, "--output", str(compile_dir)] + for config in configs: + command.extend(("--config", config)) + for value in workload: + command.extend(("--workload", str(value))) + compile_result = run(command, env=os.environ.copy()) + (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stdout.txt").write_text( + compile_result.stdout, encoding="utf-8") + (schedule_dir / f"{root}-{schedule_unique_requirement_count:03d}.stderr.txt").write_text( + compile_result.stderr, encoding="utf-8") + results.append({"case": f"schedule:{schedule_path}", "root": root, "workload": workload, + "configs": configs, "artifact_dir": str(compile_dir), + "status": "ok" if compile_result.returncode == 0 else "failed"}) summary = { "schema": "ggml-hrx-qwen-compile-report-v1", "target": args.target, "hsa_override_gfx_version": override, "hrx_revision": revision, - "corpus_digest": manifest["corpus_sha256"], - "recipe_digest": manifest["build_bazel_sha256"], + "corpus_digest": manifest.get("corpus_sha256", ""), + "recipe_digest": manifest.get("build_bazel_sha256", ""), "plan_case_count": len(manifest["plan_cases"]), "planned_invocation_count": planned_invocation_count, "resolved_invocation_count": resolved_invocation_count, "schedule_requirement_count": schedule_requirement_count, "schedule_unique_requirement_count": schedule_unique_requirement_count, "compile_count": len(results), - "failed_count": sum(item["status"] != "ok" for item in results), + "skipped_incompatible_target_count": sum( + item["status"] == "skipped-incompatible-target" for item in results + ), + "failed_count": sum(item["status"] == "failed" for item in results), "results": results, } (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") diff --git a/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py b/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py index 422b210b1830..4c5826126847 100644 --- a/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py +++ b/ggml/src/ggml-hrx/tools/generate_kernel_corpus.py @@ -107,8 +107,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--source-output", type=pathlib.Path, required=True) parser.add_argument("--corpus-output", type=pathlib.Path, required=True) parser.add_argument("--catalog-output", type=pathlib.Path, required=True) - parser.add_argument("--manifest", type=pathlib.Path, required=True) - parser.add_argument("--corpus-dir", type=pathlib.Path, required=True) + parser.add_argument("--manifest", type=pathlib.Path, action="append", required=True) + parser.add_argument("--corpus-dir", type=pathlib.Path, action="append", required=True) parser.add_argument("--source-format", choices=("text", "binary"), default="text") parser.add_argument("--loom-link", type=pathlib.Path) parser.add_argument("--loom-format", type=pathlib.Path) @@ -148,7 +148,7 @@ def convert_source_to_bytecode(source_path: pathlib.Path, loom_link: pathlib.Pat require_tool([ str(loom_link), "--verify=false", - "--mode=archive", + "--mode=merge", "--strip-check", "--to=text", f"--output={stripped}", @@ -165,7 +165,7 @@ def convert_source_to_bytecode(source_path: pathlib.Path, loom_link: pathlib.Pat require_tool([ str(loom_link), "--verify=false", - "--mode=archive", + "--mode=merge", "--strip-check", "--to=bc", f"--output={bytecode}", @@ -292,14 +292,147 @@ def collect_sources(manifest: dict) -> Tuple[List[str], Dict[str, List[str]]]: return sorted(all_sources), source_dependencies -def manifest_file_digests(manifest: dict) -> Dict[str, str]: - return {file["path"]: file["sha256"] for file in manifest.get("files", [])} - - def sha256(data: bytes) -> str: return hashlib.sha256(data).hexdigest() +def json_digest(value: object) -> str: + return sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")) + + +def complete_manifest_source_metadata(manifest: dict, source_paths: Dict[str, pathlib.Path]) -> dict: + completed = dict(manifest) + files = [] + for file in manifest.get("files", []): + path = str(file["path"]) + file_record = {key: value for key, value in file.items() if key not in ("owner", "sha256", "size")} + data = read_bytes(source_paths[path]) + file_record["sha256"] = sha256(data) + files.append(file_record) + completed["files"] = files + corpus_digest_input = { + "files": [{"path": file["path"], "sha256": file["sha256"]} for file in files], + "exports": [ + { + "family": export.get("family", DEFAULT_KERNEL_FAMILY), + "name": export["name"], + "source": export["source"], + "symbol": export["symbol"], + "target_selector": export.get("target_selector", ""), + "compile_recipe": export["compile_recipe"], + } + for export in completed.get("exports", []) + ], + } + recipe_digest_input = { + "link_modules": completed.get("link_modules", []), + "plan_cases": completed.get("plan_cases", []), + } + completed["corpus_sha256"] = json_digest(corpus_digest_input) + completed["build_bazel_sha256"] = json_digest(recipe_digest_input) + return completed + + +def merge_unique_text(values: Iterable[str]) -> str: + unique = list(dict.fromkeys(value for value in values if value)) + return ";".join(unique) + + +def require_unique_export_variants(exports: List[dict]) -> None: + variants: Dict[Tuple[str, str, str], dict] = {} + for export in exports: + key = ( + str(export.get("family", DEFAULT_KERNEL_FAMILY)), + str(export["name"]), + str(export.get("target_selector", "")), + ) + previous = variants.get(key) + if previous is not None: + family, name, target = key + label = target or "default" + raise RuntimeError(f"kernel corpus repeats target variant {family}:{name}@{label}") + variants[key] = export + + +def merge_manifests(manifest_paths: List[pathlib.Path], corpus_dirs: List[pathlib.Path]) -> Tuple[dict, Dict[str, pathlib.Path]]: + if len(manifest_paths) != len(corpus_dirs): + raise RuntimeError("--manifest and --corpus-dir must be provided the same number of times") + if not manifest_paths: + raise RuntimeError("at least one --manifest is required") + + loaded = [(manifest_path, corpus_dir, json.loads(read_text(manifest_path))) + for manifest_path, corpus_dir in zip(manifest_paths, corpus_dirs)] + if len(loaded) == 1: + manifest_path, corpus_dir, manifest = loaded[0] + source_paths = {file["path"]: corpus_dir / file["path"] for file in manifest.get("files", [])} + return complete_manifest_source_metadata(manifest, source_paths), source_paths + + files = [] + file_digests: Dict[str, str] = {} + source_paths: Dict[str, pathlib.Path] = {} + exports = [] + link_modules = [] + plan_cases = [] + metadata = [] + + for manifest_path, corpus_dir, manifest in loaded: + manifest_source_paths = {file["path"]: corpus_dir / file["path"] for file in manifest.get("files", [])} + manifest = complete_manifest_source_metadata(manifest, manifest_source_paths) + metadata.append({ + "manifest": str(manifest_path), + "upstream_revision": manifest.get("upstream_revision", ""), + "corpus_sha256": manifest.get("corpus_sha256", ""), + "build_bazel_sha256": manifest.get("build_bazel_sha256", ""), + }) + for file in manifest.get("files", []): + path = str(file["path"]) + digest = str(file["sha256"]) + previous = file_digests.get(path) + if previous is not None: + if previous != digest: + raise RuntimeError(f"manifest file digest conflict for {path}: {previous} vs {digest}") + continue + file_digests[path] = digest + source_paths[path] = corpus_dir / path + files.append(dict(file)) + exports.extend(dict(export) for export in manifest.get("exports", [])) + link_modules.extend(dict(module) for module in manifest.get("link_modules", [])) + plan_cases.extend(dict(case) for case in manifest.get("plan_cases", [])) + + require_unique_export_variants(exports) + + corpus_digest_input = { + "files": [{"path": file["path"], "sha256": file["sha256"]} for file in files], + "exports": [ + { + "family": export.get("family", DEFAULT_KERNEL_FAMILY), + "name": export["name"], + "source": export["source"], + "symbol": export["symbol"], + "target_selector": export.get("target_selector", ""), + "compile_recipe": export["compile_recipe"], + } + for export in exports + ], + } + recipe_digest_input = { + "manifests": metadata, + "link_modules": link_modules, + "plan_cases": plan_cases, + } + + return { + "schema": "ggml-hrx-kernel-corpus-v2", + "upstream_revision": merge_unique_text(str(manifest.get("upstream_revision", "")) for _, _, manifest in loaded), + "corpus_sha256": json_digest(corpus_digest_input), + "build_bazel_sha256": json_digest(recipe_digest_input), + "files": files, + "exports": exports, + "link_modules": link_modules, + "plan_cases": plan_cases, + }, source_paths + + def kernel_catalog_id(family: str, name: str) -> int: hash_value = 1469598103934665603 for byte in family.encode("utf-8"): @@ -313,8 +446,7 @@ def kernel_catalog_id(family: str, name: str) -> int: return hash_value -def generate_corpus_records(manifest: dict, source_records: Dict[str, str]) -> Tuple[str, str, int]: - digests = manifest_file_digests(manifest) +def generate_corpus_records(manifest: dict, source_records: Dict[str, str], source_digests: Dict[str, str]) -> Tuple[str, str, int]: arrays = [] records = [] exports = manifest.get("exports", []) @@ -360,7 +492,7 @@ def generate_corpus_records(manifest: dict, source_records: Dict[str, str]) -> T target_selector=cpp_string(export.get("target_selector", "")), source=cpp_string(export["source"]), dependencies=dependencies_span, - source_digest=cpp_string(digests[export["source"]]), + source_digest=cpp_string(source_digests[export["source"]]), scalar_parameters=scalar_span, bindings=bindings_span, workload_parameters=workload_span, @@ -396,13 +528,12 @@ def generate_catalog_verifier(manifest: dict) -> str: ) -def generate_includes(args: argparse.Namespace, manifest: dict) -> Tuple[str, str, str, List[pathlib.Path], int]: +def generate_includes(args: argparse.Namespace, manifest: dict, source_paths: Dict[str, pathlib.Path]) -> Tuple[str, str, str, List[pathlib.Path], int]: if args.source_format == "binary" and (args.loom_link is None or args.loom_format is None): raise RuntimeError("binary source format requires --loom-link and --loom-format") - corpus_dir = args.corpus_dir sources, source_dependencies = collect_sources(manifest) - digests = manifest_file_digests(manifest) + source_digests = {str(file["path"]): str(file["sha256"]) for file in manifest.get("files", [])} source_bytes: Dict[str, bytes] = {} source_format = "KERNEL_SOURCE_FORMAT_BINARY" if args.source_format == "binary" else "KERNEL_SOURCE_FORMAT_TEXT" input_files: List[pathlib.Path] = [] @@ -414,14 +545,12 @@ def generate_includes(args: argparse.Namespace, manifest: dict) -> Tuple[str, st raise RuntimeError(f"export primary source is not embedded: {primary}") for source in sources: - if source not in digests: + if source not in source_digests: raise RuntimeError(f"embedded source is missing from manifest file table: {source}") - path = corpus_dir / source + path = source_paths[source] input_files.append(path) data = read_bytes(path) - digest = sha256(data) - if digest != digests[source]: - raise RuntimeError(f"manifest digest mismatch for {source}: got {digest}, expected {digests[source]}") + source_digests[source] = sha256(data) if args.source_format == "binary": data = convert_source_to_bytecode(path, args.loom_link, args.loom_format) source_bytes[source] = data @@ -476,7 +605,7 @@ def generate_includes(args: argparse.Namespace, manifest: dict) -> Tuple[str, st ) lookup_entries.append(f" {{ {cpp_string(source)}, &{record} }},") - kernel_arrays, kernel_records, kernel_count = generate_corpus_records(manifest, source_record_symbols) + kernel_arrays, kernel_records, kernel_count = generate_corpus_records(manifest, source_record_symbols, source_digests) return ( SOURCE_DATA_TEMPLATE.format( source_arrays="\n".join(source_arrays), @@ -508,8 +637,9 @@ def write_depfile(path: pathlib.Path, outputs: List[pathlib.Path], inputs: List[ def main() -> int: args = parse_args() try: - manifest = json.loads(read_text(args.manifest)) - source_include, corpus_include, catalog_include, input_files, byte_count = generate_includes(args, manifest) + manifest, source_paths = merge_manifests(args.manifest, args.corpus_dir) + source_include, corpus_include, catalog_include, input_files, byte_count = generate_includes( + args, manifest, source_paths) args.source_output.parent.mkdir(parents=True, exist_ok=True) args.corpus_output.parent.mkdir(parents=True, exist_ok=True) args.catalog_output.parent.mkdir(parents=True, exist_ok=True) @@ -519,7 +649,7 @@ def main() -> int: if args.depfile is not None: args.depfile.parent.mkdir(parents=True, exist_ok=True) write_depfile(args.depfile, [args.source_output, args.corpus_output, args.catalog_output], - [args.manifest, *input_files]) + [*args.manifest, *input_files]) except Exception as exc: print(f"generate_kernel_corpus.py: {exc}", file=sys.stderr) return 1 diff --git a/ggml/src/ggml-hrx/tools/sync_qwen_kernel_corpus.py b/ggml/src/ggml-hrx/tools/sync_qwen_kernel_corpus.py deleted file mode 100755 index 5d4d9705c47d..000000000000 --- a/ggml/src/ggml-hrx/tools/sync_qwen_kernel_corpus.py +++ /dev/null @@ -1,596 +0,0 @@ -#!/usr/bin/env python3 -"""Mirrors the bounded Qwen MoE Loom corpus with reproducible provenance.""" - -from __future__ import annotations - -import argparse -import ast -import hashlib -import json -import pathlib -import re -import shutil -import subprocess -import sys -import tempfile - - -SOURCE_SUBDIR = pathlib.Path("experimental/qwen_moe/kernels") -QWEN_ENDPOINT_SOURCE_SUBDIR = pathlib.Path("experimental/qwen/kernels") -CORPUS_FILES = ( - "ggml/linear_q6k_f32.loom", - "ggml/linear_q6k_q8_1_x4.loom", - "ggml/quantize_q8_1_x4.loom", - "qwen3_moe/attention_postprocess_f32_f16.loom", - "qwen3_moe/attention_prepare_quantized.loom", - "qwen3_moe/attention_qkv_postprocess_fused.loom", - "qwen3_moe/attention_qkv_quantized.loom", - "qwen3_moe/attention_qkv_same_format_prefill.loom", - "qwen3_moe/batched_decode_expert_dispatch.loom", - "qwen3_moe/batched_decode_gate_up_q4k.loom", - "qwen3_moe/dense_linear_quantized_f16_wmma.loom", - "qwen3_moe/expert_table_partition_fused.loom", - "qwen3_moe/flash_attention_decode_f32_f16_wmma.loom", - "qwen3_moe/flash_attention_decode_q128_f32_f16_wmma.loom", - "qwen3_moe/flash_attention_decode_split_f32_f16_wmma.loom", - "qwen3_moe/flash_attention_decode_split_next_q8_test.loom", - "qwen3_moe/flash_attention_f32_f16_wmma.loom", - "qwen3_moe/model_config.loom", - "qwen3_moe/routed_down_q4k.loom", - "qwen3_moe/routed_down_q6k.loom", - "qwen3_moe/routed_down_next_q8.loom", - "qwen3_moe/routed_down_quantized_f16_wmma.loom", - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_f32.loom", - "qwen3_moe/routed_down_weighted_reduce_next_rmsnorm_q8_1_x4.loom", - "qwen3_moe/routed_gate_up_swiglu_q4k.loom", - "qwen3_moe/routed_linear_q4k_f16_wmma.loom", - "qwen3_moe/router_projection_f32.loom", - "qwen3_moe/router_projection_top8_fused_f32.loom", - "qwen3_moe/router_top8_f32.loom", -) -QWEN_ENDPOINT_FILES = ( - ("token_embedding_q4k.loom", "qwen_owned/token_embedding_q4k.loom"), - ("attention_metadata.loom", "qwen_owned/attention_metadata.loom"), -) - -# These integration kernels are deliberately owned by the llama.cpp HRX -# backend. They are not attributed to the pinned qwen_moe corpus or its BUILD -# recipes. -OWNED_KERNEL_DIR = pathlib.Path(__file__).resolve().parent.parent / "kernel-corpus" / "kernels" -OWNED_FILES = ( - "qwen_owned/token_embedding_bringup_workaround.loom", - "qwen_owned/attention_state_initialize.loom", - "qwen_owned/attention_metadata_bringup_workaround.loom", - "hrx_owned/gather_add_f32.loom", - "hrx_owned/add_f32.loom", -) - -KERNEL_RE = re.compile( - r"kernel\.def(?P(?:\s+(?:target\([^)]*\)|export\(\"[^\"]+\"\)))*)" - r"\s+@(?P[A-Za-z0-9_]+)" - r"\((?P.*?)\)\s*\{.*?\}\s*launch\((?P.*?)\)" - r"(?:\s+where\s+\[[^\]]*\])?\s*\{", - re.DOTALL, -) -ARG_RE = re.compile(r"%(?P[A-Za-z0-9_]+)\s*:\s*(?P[A-Za-z0-9<>?]+)") -TARGET_MODIFIER_RE = re.compile(r"target\(@(?P[A-Za-z0-9_]+)\)") -EXPORT_MODIFIER_RE = re.compile(r"export\(\"(?P[^\"]+)\"\)") -AMDGPU_TARGET_RE = re.compile( - r"amdgpu\.target<(?P[A-Za-z0-9_.-]+)>\s+@(?P[A-Za-z0-9_]+)" -) - - -def binding_access(symbol: str, name: str) -> str: - """Authoritative launch ABI access contract; no name inference at runtime.""" - if symbol == "qwen_attention_context_base_capture": - return "read" if name == "positions" else "write" - if symbol == "qwen_attention_decode_state_initialize": - return "read" if name == "positions" else "write" - if symbol == "qwen_attention_metadata_bringup_workaround" and name != "control": - return "read_write" - if symbol == "qwen_attention_metadata" and name != "control": - return "read_write" - if symbol == "qwen_decode_attention_metadata" and name != "control": - return "read_write" - if symbol == "qwen3_moe_router_top8_f32" and name in ("route_ids", "route_weights"): - return "write" - if symbol == "qwen3_moe_router_projection_top8_fused_decode_f32" and name in ( - "logits", "completion_counter", "route_ids", "route_weights"): - return "read_write" - if symbol == "qwen3_moe_attention_qkv_postprocess_fused_decode" and name in ( - "query_output_raw", "key_output_raw", "value_output_raw", - "query_output", "key_cache", "value_cache", "completion_counters"): - return "read_write" - if symbol == "qwen3_moe_build_expert_table" and name == "expert_table": - return "write" - if symbol == "qwen3_moe_build_expert_partition_table" and name == "partition_table": - return "write" - if symbol == "qwen3_moe_build_expert_table_partition_prefill_512" and name in ("expert_table", "partition_table"): - return "write" - if symbol == "qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32": - if name == "hidden_state": - return "read_write" - if name == "next_projection_input": - return "write" - if symbol == "ggml_q8_1_x4_inspect_one_group" and name != "packed": - return "write" - if name in ("output", "query_output", "key_output", "value_output", "normalized_output", "q8_output", - "next_q8_output", "key_cache", "value_cache", "partial_max", "partial_sum", "partial_output", - "completion_counter", "completion_counters"): - return "read_write" - return "read" - - -def starlark_calls(text: str, function: str) -> list[str]: - """Extracts the literal-only calls used by the pinned kernel BUILD file.""" - result: list[str] = [] - marker = function + "(" - cursor = 0 - while (start := text.find(marker, cursor)) != -1: - index = start + len(marker) - depth = 1 - quote: str | None = None - escaped = False - while index < len(text) and depth: - character = text[index] - if quote: - if escaped: - escaped = False - elif character == "\\": - escaped = True - elif character == quote: - quote = None - elif character in "\"'": - quote = character - elif character == "(": - depth += 1 - elif character == ")": - depth -= 1 - index += 1 - if depth: - raise RuntimeError(f"unterminated {function} call in BUILD.bazel") - result.append(text[start + len(marker) : index - 1]) - cursor = index - return result - - -def literal_assignment(call: str, name: str, default: object = None) -> object: - match = re.search(rf"(?:^|\n)\s*{re.escape(name)}\s*=\s*(\[[\s\S]*?\]|\"[^\"]*\")\s*,", call) - return default if match is None else ast.literal_eval(match.group(1)) - - -def parse_build_recipes(text: str) -> tuple[list[dict[str, object]], list[dict[str, object]]]: - modules: list[dict[str, object]] = [] - for call in starlark_calls(text, "loom_link_module"): - modules.append({ - "name": literal_assignment(call, "name"), - "srcs": literal_assignment(call, "srcs", []), - "libraries": literal_assignment(call, "libraries", []), - }) - module_names = {str(module["name"]) for module in modules} - cases: list[dict[str, object]] = [] - for call in starlark_calls(text, "iree_executable_test"): - name = literal_assignment(call, "name") - if not isinstance(name, str) or not name.endswith("_plan_test"): - continue - args = literal_assignment(call, "args", []) - data = literal_assignment(call, "data", []) - linked = [item[1:] for item in data if isinstance(item, str) and item.startswith(":") and item[1:] in module_names] - direct_sources = [item for item in data if isinstance(item, str) and item.endswith(".loom")] - if len(linked) + len(direct_sources) != 1: - raise RuntimeError(f"plan test {name} does not name exactly one linked module or direct Loom source") - case = { - "name": name, - "args": args, - } - if linked: - case["link_module"] = linked[0] - else: - case["source"] = direct_sources[0] - cases.append(case) - if not modules or not cases: - raise RuntimeError("BUILD.bazel contains no pinned Loom link/plan recipes") - return modules, cases - - -def git(repo: pathlib.Path, *args: str) -> str: - return subprocess.check_output(["git", "-C", str(repo), *args], text=True).strip() - - -def optional_git(repo: pathlib.Path, *args: str) -> str | None: - result = subprocess.run( - ["git", "-C", str(repo), *args], text=True, - stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, check=False) - value = result.stdout.strip() - return value if result.returncode == 0 and value else None - - -def normalize_repository_url(url: str) -> str: - """Normalizes GitHub transport spelling without changing repository identity.""" - match = re.fullmatch(r"(?:ssh://)?git@github\.com[:/](?P.+)", url) - if match: - return f"https://github.com/{match.group('path')}" - return url - - -def upstream_repository(repo: pathlib.Path) -> str: - """Returns provenance without imposing a local Git remote name.""" - remotes: list[str] = [] - branch = optional_git(repo, "symbolic-ref", "--quiet", "--short", "HEAD") - if branch: - branch_remote = optional_git(repo, "config", "--get", f"branch.{branch}.remote") - if branch_remote and branch_remote != ".": - remotes.append(branch_remote) - remotes.append("origin") - remotes.extend(git(repo, "remote").splitlines()) - for remote in dict.fromkeys(remotes): - url = optional_git(repo, "config", "--get", f"remote.{remote}.url") - if url: - return normalize_repository_url(url) - raise RuntimeError("HRX source tree has no repository remote for provenance") - - -def sha256(data: bytes) -> str: - return hashlib.sha256(data).hexdigest() - - -def parse_exports(text: str, source: str) -> list[dict[str, object]]: - exports: list[dict[str, object]] = [] - for match in KERNEL_RE.finditer(text): - modifiers = match.group("modifiers") - target_match = TARGET_MODIFIER_RE.search(modifiers) - export_match = EXPORT_MODIFIER_RE.search(modifiers) - workload = [item.groupdict() for item in ARG_RE.finditer(match.group("workload"))] - launch = [item.groupdict() for item in ARG_RE.finditer(match.group("launch"))] - bindings = [item["name"] for item in launch if item["type"] == "buffer"] - exports.append( - { - "name": export_match.group("name") if export_match else match.group("symbol"), - "symbol": match.group("symbol"), - "target_symbol": target_match.group("symbol") if target_match else "", - "source": source, - "workload_parameters": workload, - "launch_parameters": [item for item in launch if item["type"] != "buffer"], - "bindings": bindings, - "binding_access": [binding_access(match.group("symbol"), name) for name in bindings], - } - ) - return exports - - -def parse_amdgpu_targets(text: str) -> dict[str, str]: - result: dict[str, str] = {} - for match in AMDGPU_TARGET_RE.finditer(text): - symbol = match.group("symbol") - selector = match.group("selector") - if symbol in result and result[symbol] != selector: - raise RuntimeError( - f"AMDGPU target @{symbol} is declared as both {result[symbol]} and {selector}") - result[symbol] = selector - return result - - -def merge_amdgpu_targets(target_selectors: dict[str, str], additions: dict[str, str]) -> None: - for symbol, selector in additions.items(): - if symbol in target_selectors and target_selectors[symbol] != selector: - raise RuntimeError( - f"AMDGPU target @{symbol} is declared as both {target_selectors[symbol]} and {selector}") - target_selectors[symbol] = selector - - -def resolve_export_variants(exports: list[dict[str, object]], target_selectors: dict[str, str]) -> None: - groups: dict[str, list[dict[str, object]]] = {} - for item in exports: - groups.setdefault(str(item["name"]), []).append(item) - - for name, variants in groups.items(): - selectors: set[str] = set() - for item in variants: - target_symbol = str(item.pop("target_symbol")) - selector = "" - if target_symbol: - if target_symbol not in target_selectors: - raise RuntimeError(f"kernel export {name} references unknown target @{target_symbol}") - selector = target_selectors[target_symbol] - if selector.endswith("-generic"): - selector = "" - if selector in selectors: - label = selector or "default" - raise RuntimeError(f"kernel export {name} repeats target variant {label}") - selectors.add(selector) - item["target_selector"] = selector - - -def construct(source_root: pathlib.Path, destination: pathlib.Path, expected_revision: str | None) -> None: - revision = git(source_root, "rev-parse", "HEAD") - if expected_revision and revision != expected_revision: - raise RuntimeError(f"HRX revision {revision} does not match expected {expected_revision}") - if git(source_root, "status", "--porcelain"): - raise RuntimeError("refusing to mirror a dirty HRX source tree") - - source_directory = source_root / SOURCE_SUBDIR - build_data = (source_directory / "BUILD.bazel").read_bytes() - all_link_modules, all_plan_cases = parse_build_recipes(build_data.decode("utf-8")) - - modules_by_name = {str(item["name"]): item for item in all_link_modules} - direct_plan_sources = { - str(item["source"]) - for item in all_plan_cases - if "source" in item - } - - def module_files(name: str) -> list[str]: - module = modules_by_name[name] - result = list(module["srcs"]) - for library in module["libraries"]: - if str(library).startswith(":"): - result.extend(module_files(str(library)[1:])) - else: - result.append(str(library)) - return list(dict.fromkeys(result)) - - def compile_recipe(source: str) -> dict[str, object]: - direct = [str(item["name"]) for item in all_link_modules if source in item["srcs"]] - indirect = [str(item["name"]) for item in all_link_modules if source in item["libraries"]] - if not direct and (source in direct_plan_sources or not indirect): - return {"mode": "direct", "primary_sources": [source], "library_sources": []} - module_name = (direct or indirect)[0] - module = modules_by_name[module_name] - files = module_files(module_name) - return { - "mode": "archive", - "link_module": module_name, - "primary_sources": list(module["srcs"]), - "library_sources": [item for item in files if item not in module["srcs"]], - } - file_rows: list[dict[str, object]] = [] - exports: list[dict[str, object]] = [] - target_selectors: dict[str, str] = {} - upstream_aggregate = hashlib.sha256() - for relative_text in CORPUS_FILES: - relative = pathlib.Path(relative_text) - source = source_directory / relative - if not source.is_file(): - raise RuntimeError(f"missing required corpus source: {source}") - data = source.read_bytes() - digest = sha256(data) - upstream_aggregate.update(relative_text.encode()) - upstream_aggregate.update(b"\0") - upstream_aggregate.update(bytes.fromhex(digest)) - target = destination / relative - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(data) - file_rows.append({"path": relative_text, "sha256": digest, "size": len(data)}) - source_text = data.decode("utf-8") - exports.extend(parse_exports(source_text, relative_text)) - merge_amdgpu_targets(target_selectors, parse_amdgpu_targets(source_text)) - - endpoint_source_directory = source_root / QWEN_ENDPOINT_SOURCE_SUBDIR - for source_text_name, local_text_name in QWEN_ENDPOINT_FILES: - source = endpoint_source_directory / source_text_name - if not source.is_file(): - raise RuntimeError(f"missing required Qwen endpoint source: {source}") - data = source.read_bytes() - digest = sha256(data) - provenance_path = f"{QWEN_ENDPOINT_SOURCE_SUBDIR.as_posix()}/{source_text_name}" - upstream_aggregate.update(provenance_path.encode()) - upstream_aggregate.update(b"\0") - upstream_aggregate.update(bytes.fromhex(digest)) - relative_text = f"../{local_text_name}" - target = destination.parent / local_text_name - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(data) - file_rows.append({ - "path": relative_text, - "sha256": digest, - "size": len(data), - "upstream_path": provenance_path, - }) - source_text = data.decode("utf-8") - exports.extend(parse_exports(source_text, relative_text)) - merge_amdgpu_targets(target_selectors, parse_amdgpu_targets(source_text)) - - owned_aggregate = hashlib.sha256() - for filename in OWNED_FILES: - source = OWNED_KERNEL_DIR / filename - if not source.is_file(): - raise RuntimeError(f"missing required backend-owned kernel source: {source}") - data = source.read_bytes() - digest = sha256(data) - relative_text = f"../{filename}" - owned_aggregate.update(filename.encode()) - owned_aggregate.update(b"\0") - owned_aggregate.update(bytes.fromhex(digest)) - file_rows.append({"path": relative_text, "sha256": digest, "size": len(data), "owner": "ggml-hrx"}) - source_text = data.decode("utf-8") - exports.extend(parse_exports(source_text, relative_text)) - merge_amdgpu_targets(target_selectors, parse_amdgpu_targets(source_text)) - - resolve_export_variants(exports, target_selectors) - - for item in exports: - recipe = compile_recipe(str(item["source"])) - item["compile_recipe"] = recipe - item["compile_dependencies"] = list(recipe["library_sources"]) - - required_modules: set[str] = set() - - def require_module(name: str) -> None: - if name in required_modules: - return - if name not in modules_by_name: - raise RuntimeError(f"selected kernel recipe references unknown link module {name}") - required_modules.add(name) - for library in modules_by_name[name]["libraries"]: - if str(library).startswith(":"): - require_module(str(library)[1:]) - - required_files = set(CORPUS_FILES) - for item in exports: - if str(item["source"]).startswith("../"): - continue - recipe = item["compile_recipe"] - link_module = str(recipe.get("link_module", "")) - if link_module: - require_module(link_module) - required_files.update(str(path) for path in recipe["primary_sources"]) - required_files.update(str(path) for path in recipe["library_sources"]) - - mirrored_files = set(CORPUS_FILES) - for relative_text in sorted(required_files - mirrored_files): - relative = pathlib.Path(relative_text) - if relative.is_absolute() or ".." in relative.parts: - raise RuntimeError(f"kernel recipe escapes the source corpus: {relative_text}") - source = source_directory / relative - if not source.is_file(): - raise RuntimeError(f"missing required kernel dependency: {source}") - data = source.read_bytes() - digest = sha256(data) - upstream_aggregate.update(relative_text.encode()) - upstream_aggregate.update(b"\0") - upstream_aggregate.update(bytes.fromhex(digest)) - target = destination / relative - target.parent.mkdir(parents=True, exist_ok=True) - target.write_bytes(data) - file_rows.append({"path": relative_text, "sha256": digest, "size": len(data)}) - - link_modules = [ - module for module in all_link_modules - if str(module["name"]) in required_modules - ] - plan_cases = [ - case for case in all_plan_cases - if (str(case.get("link_module", "")) in required_modules or - str(case.get("source", "")) in required_files) - ] - plan_cases.extend([ - { - "name": "owned_token_embedding_decode_plan_test", - "args": ["$(location ../qwen_owned/token_embedding_q4k.loom)", - "--benchmark=@qwen_token_embedding_q4k_decode", "--dry-run", - "--output-format=jsonl", "--sample-compilation=per_sample"], - "source": "../qwen_owned/token_embedding_q4k.loom", - }, - { - "name": "owned_token_embedding_prefill_plan_test", - "args": ["$(location ../qwen_owned/token_embedding_q4k.loom)", - "--benchmark=@qwen_token_embedding_q4k_prefill_512", "--dry-run", - "--output-format=jsonl", "--sample-compilation=per_sample"], - "source": "../qwen_owned/token_embedding_q4k.loom", - }, - { - "name": "owned_attention_context_base_capture_plan_test", - "args": ["$(location ../qwen_owned/attention_state_initialize.loom)", - "--benchmark=@qwen_attention_context_base_capture_benchmark", "--dry-run", - "--output-format=jsonl", "--sample-compilation=per_sample"], - "source": "../qwen_owned/attention_state_initialize.loom", - "owner": "ggml-hrx", - }, - { - "name": "owned_attention_decode_state_initialize_plan_test", - "args": ["$(location ../qwen_owned/attention_state_initialize.loom)", - "--benchmark=@qwen_attention_decode_state_initialize_benchmark", "--dry-run", - "--output-format=jsonl", "--sample-compilation=per_sample"], - "source": "../qwen_owned/attention_state_initialize.loom", - "owner": "ggml-hrx", - }, - { - "name": "owned_attention_metadata_prefill_plan_test", - "args": ["$(location ../qwen_owned/attention_metadata.loom)", - "--benchmark=@qwen_attention_metadata_prefill_512", "--dry-run", - "--output-format=jsonl", "--sample-compilation=per_sample"], - "source": "../qwen_owned/attention_metadata.loom", - }, - { - "name": "owned_gather_add_plan_test", - "args": ["$(location ../hrx_owned/gather_add_f32.loom)", - "--benchmark=@ggml_gather_add_noncontiguous", "--dry-run", - "--output-format=jsonl", "--sample-compilation=per_sample"], - "source": "../hrx_owned/gather_add_f32.loom", - "owner": "ggml-hrx", - }, - ]) - - upstream_digest = upstream_aggregate.hexdigest() - owned_digest = owned_aggregate.hexdigest() - combined_aggregate = hashlib.sha256() - combined_aggregate.update(bytes.fromhex(upstream_digest)) - combined_aggregate.update(bytes.fromhex(owned_digest)) - - manifest = { - "schema": "ggml-hrx-qwen-kernel-corpus-v2", - "upstream_repository": upstream_repository(source_root), - "upstream_revision": revision, - "source_subdirectory": SOURCE_SUBDIR.as_posix(), - "qwen_endpoint_source_subdirectory": QWEN_ENDPOINT_SOURCE_SUBDIR.as_posix(), - "corpus_sha256": combined_aggregate.hexdigest(), - "upstream_corpus_sha256": upstream_digest, - "owned_corpus_sha256": owned_digest, - "build_bazel_sha256": sha256(build_data), - "files": file_rows, - "exports": sorted(exports, key=lambda item: (str(item["symbol"]), str(item["source"]))), - "link_modules": link_modules, - "plan_cases": plan_cases, - } - (destination / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") - - -def trees_equal(lhs: pathlib.Path, rhs: pathlib.Path) -> bool: - lhs_files = sorted(path.relative_to(lhs) for path in lhs.rglob("*") if path.is_file()) - rhs_files = sorted(path.relative_to(rhs) for path in rhs.rglob("*") if path.is_file()) - return lhs_files == rhs_files and all((lhs / path).read_bytes() == (rhs / path).read_bytes() for path in lhs_files) - - -def endpoint_files_equal(generated_qwen_moe: pathlib.Path, destination_qwen_moe: pathlib.Path) -> bool: - generated_kernel_root = generated_qwen_moe.parent - destination_kernel_root = destination_qwen_moe.parent - for _, local_text_name in QWEN_ENDPOINT_FILES: - relative = pathlib.Path(local_text_name) - generated = generated_kernel_root / relative - destination = destination_kernel_root / relative - if not destination.is_file() or generated.read_bytes() != destination.read_bytes(): - return False - return True - - -def copy_endpoint_files(generated_qwen_moe: pathlib.Path, destination_qwen_moe: pathlib.Path) -> None: - generated_kernel_root = generated_qwen_moe.parent - destination_kernel_root = destination_qwen_moe.parent - for _, local_text_name in QWEN_ENDPOINT_FILES: - relative = pathlib.Path(local_text_name) - source = generated_kernel_root / relative - target = destination_kernel_root / relative - target.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(source, target) - - -def main() -> int: - parser = argparse.ArgumentParser() - parser.add_argument("--hrx-source", type=pathlib.Path, required=True) - parser.add_argument("--destination", type=pathlib.Path, required=True) - parser.add_argument("--expect-revision") - parser.add_argument("--check", action="store_true") - args = parser.parse_args() - - with tempfile.TemporaryDirectory(prefix="hrx-qwen-corpus-") as temporary: - generated = pathlib.Path(temporary) / args.destination.name - construct(args.hrx_source.resolve(), generated, args.expect_revision) - if args.check: - if (not args.destination.is_dir() or not trees_equal(generated, args.destination) or - not endpoint_files_equal(generated, args.destination)): - print("mirrored Qwen kernel corpus is stale", file=sys.stderr) - return 1 - return 0 - args.destination.mkdir(parents=True, exist_ok=True) - for child in args.destination.iterdir(): - if child.is_dir(): - shutil.rmtree(child) - else: - child.unlink() - shutil.copytree(generated, args.destination, dirs_exist_ok=True) - copy_endpoint_files(generated, args.destination) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 1320f1255f66..cfe17209aef0 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -5,9 +5,12 @@ #include "dispatch/command-program-resolver.h" #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" +#include "ggml-cpu.h" #include "ggml-hrx.h" #include "ggml-impl.h" #include "ggml.h" @@ -22,6 +25,7 @@ #include "runtime/graph-program-cache.h" #include "runtime/graph-replay.h" #include "runtime/loom-kernel-jit.h" +#include "testing_suite.h" #include #include @@ -29,9 +33,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -66,7 +72,14 @@ static bool contains_value_id(const std::vector & ids, ggml: } static bool command_program_verifies(const ggml::hrx::CommandProgram & program) { - return ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151").valid(); + const ggml::hrx::VerificationResult result = + ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + if (!result.valid()) { + for (const std::string & error : result.status.errors()) { + std::fprintf(stderr, "%s\n", error.c_str()); + } + } + return result.valid(); } static bool async_jit_expected_from_environment() { @@ -99,6 +112,61 @@ static bool string_contains(const std::string & value, const char * text) { return value.find(text) != std::string::npos; } +static std::string read_text_file(const std::filesystem::path & path) { + std::ifstream input(path, std::ios::binary); + REQUIRE(input.good()); + return { std::istreambuf_iterator(input), std::istreambuf_iterator() }; +} + +static std::vector list_directories(const std::filesystem::path & path) { + std::vector directories; + if (!std::filesystem::exists(path)) { + return directories; + } + for (const std::filesystem::directory_entry & entry : std::filesystem::directory_iterator(path)) { + if (entry.is_directory()) { + directories.push_back(entry.path()); + } + } + std::sort(directories.begin(), directories.end()); + return directories; +} + +struct EnvironmentVariableGuard { + const char * name; + bool had_value = false; + std::string value; + + explicit EnvironmentVariableGuard(const char * name) : name(name) { + const char * current = std::getenv(name); + if (current != nullptr) { + had_value = true; + value = current; + } + } + + ~EnvironmentVariableGuard() { + if (had_value) { + setenv(name, value.c_str(), 1); + } else { + unsetenv(name); + } + } + + void set(const std::filesystem::path & path) const { setenv(name, path.string().c_str(), 1); } + + void unset() const { unsetenv(name); } +}; + +static std::filesystem::path fresh_test_directory(const char * name) { + static uint64_t sequence = 0; + std::filesystem::path path = std::filesystem::temp_directory_path() / + ("hrx-backend-test-" + std::string(name) + "-" + std::to_string(sequence++)); + std::filesystem::remove_all(path); + std::filesystem::create_directories(path); + return path; +} + static void require_hrx_status(hrx_status_t status) { if (ggml::hrx::ErrorResult error = ggml::hrx::take_status(status)) { std::fprintf(stderr, "HRX status failed: %s\n", error->c_str()); @@ -123,12 +191,53 @@ static std::string kernel_name_for_id(uint64_t kernel_id) { return ggml::hrx::kernel_definition_name(*resolved.definition); } +template static void require_scheduled_command_program(ggml_cgraph * graph, Check && check) { + 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()); + + 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)); + check(imported.graph, scheduler.plan(), commands); +} + +static std::vector scheduled_kernel_names(ggml_cgraph * graph) { + 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()); + + std::vector result; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + result.push_back(kernel_name_for_id(dispatch.kernel.kernel_id)); + } + return result; +} + static void require_compile_parameter(const ggml::hrx::Dispatch & dispatch, const char * name, const std::string & value) { const auto found = dispatch.kernel.compile_parameters.find(name); REQUIRE(found != dispatch.kernel.compile_parameters.end()); - REQUIRE(found->second == value); + if (found->second != value) { + std::fprintf(stderr, "compile parameter %s expected '%s' got '%s'\n", name, value.c_str(), + found->second.c_str()); + std::abort(); + } +} + +static std::string expected_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); } static constexpr int64_t kQwenFlashHeadSize = 128; @@ -153,6 +262,54 @@ static size_t qwen_q8_1_x4_size(int64_t token_count, int64_t hidden_size) { return static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, hidden_size); } +static int64_t matmul_weight_format_config(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: + return 10; + case GGML_TYPE_Q3_K: + return 11; + case GGML_TYPE_Q4_K: + return 4; + case GGML_TYPE_Q5_K: + return 5; + case GGML_TYPE_Q6_K: + return 6; + case GGML_TYPE_Q4_0: + return 40; + case GGML_TYPE_Q4_1: + return 41; + case GGML_TYPE_Q5_0: + 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: + return 21; + case GGML_TYPE_IQ4_NL: + return 20; + case GGML_TYPE_IQ4_XS: + return 23; + case GGML_TYPE_Q8_0: + return 80; + case GGML_TYPE_Q8_1: + return 81; + case GGML_TYPE_F16: + return 16; + case GGML_TYPE_BF16: + return 30; + case GGML_TYPE_F32: + return 32; + default: + REQUIRE(false); + return 0; + } +} + static std::vector make_qwen_route_ids_iota(int64_t token_count, int64_t route_count, int64_t route_stride, @@ -249,18 +406,18 @@ static size_t qwen_routed_down_f16_output_size(int64_t token_count) { return static_cast(token_count * kQwenRouterRouteCount * kQwenMoeHiddenSize) * sizeof(ggml_fp16_t); } -static void set_qwen_flash_query_layout(ggml_tensor * tensor, int64_t head_count) { +static void set_qwen_flash_query_layout(ggml_tensor * tensor, int64_t head_count, int64_t head_size) { REQUIRE(tensor != nullptr); tensor->nb[0] = sizeof(float); - tensor->nb[1] = static_cast(head_count * kQwenFlashHeadSize) * sizeof(float); - tensor->nb[2] = static_cast(kQwenFlashHeadSize) * sizeof(float); + tensor->nb[1] = static_cast(head_count * head_size) * sizeof(float); + tensor->nb[2] = static_cast(head_size) * sizeof(float); } -static void set_qwen_flash_key_value_layout(ggml_tensor * tensor, int64_t head_count) { +static void set_qwen_flash_key_value_layout(ggml_tensor * tensor, int64_t head_count, int64_t head_size) { REQUIRE(tensor != nullptr); tensor->nb[0] = sizeof(ggml_fp16_t); - tensor->nb[1] = static_cast(head_count * kQwenFlashHeadSize) * sizeof(ggml_fp16_t); - tensor->nb[2] = static_cast(kQwenFlashHeadSize) * sizeof(ggml_fp16_t); + tensor->nb[1] = static_cast(head_count * head_size) * sizeof(ggml_fp16_t); + tensor->nb[2] = static_cast(head_size) * sizeof(ggml_fp16_t); } static ggml_tensor * build_qwen_flash_attention_graph(ggml_context * ctx, @@ -268,27 +425,29 @@ static ggml_tensor * build_qwen_flash_attention_graph(ggml_context * ctx, int64_t key_value_token_count, int64_t query_head_count, int64_t key_value_head_count, - ggml_type query_type = GGML_TYPE_F32, - ggml_type key_value_type = GGML_TYPE_F16, - bool include_mask = true, - bool include_sinks = false, - int64_t head_size = kQwenFlashHeadSize, - float scale = 1.0f / std::sqrt(128.0f), - float max_bias = 0.0f, - float logit_softcap = 0.0f) { + ggml_type query_type = GGML_TYPE_F32, + ggml_type key_value_type = GGML_TYPE_F16, + bool include_mask = true, + bool include_sinks = false, + int64_t head_size = kQwenFlashHeadSize, + float scale = 1.0f / std::sqrt(128.0f), + int64_t value_head_size = -1, + float max_bias = 0.0f, + float logit_softcap = 0.0f) { + const int64_t actual_value_head_size = value_head_size > 0 ? value_head_size : head_size; ggml_tensor * query = ggml_new_tensor_3d(ctx, query_type, head_size, query_token_count, query_head_count); ggml_tensor * key = ggml_new_tensor_3d(ctx, key_value_type, head_size, key_value_token_count, key_value_head_count); ggml_tensor * value = - ggml_new_tensor_3d(ctx, key_value_type, head_size, key_value_token_count, key_value_head_count); + ggml_new_tensor_3d(ctx, key_value_type, actual_value_head_size, key_value_token_count, key_value_head_count); REQUIRE(query != nullptr); REQUIRE(key != nullptr); REQUIRE(value != nullptr); if (query_type == GGML_TYPE_F32) { - set_qwen_flash_query_layout(query, query_head_count); + set_qwen_flash_query_layout(query, query_head_count, head_size); } if (key_value_type == GGML_TYPE_F16) { - set_qwen_flash_key_value_layout(key, key_value_head_count); - set_qwen_flash_key_value_layout(value, key_value_head_count); + set_qwen_flash_key_value_layout(key, key_value_head_count, head_size); + set_qwen_flash_key_value_layout(value, key_value_head_count, actual_value_head_size); } ggml_tensor * mask = nullptr; if (include_mask) { @@ -433,8 +592,8 @@ static void run_command_plan_metadata_checks() { { ggml::hrx::ValueId(1), ggml::hrx::ValueId(3), GGML_TYPE_F16, 16, "alternate" }, status)); REQUIRE(!status.success()); - ggml::hrx::CommandPlanMetadata bundle_plan; - ggml::hrx::Status bundle_status; + ggml::hrx::CommandPlanMetadata bundle_plan; + ggml::hrx::Status bundle_status; const ggml::hrx::CommandPlanMoeRoutingBundle bundle = { ggml::hrx::ValueId(10), ggml::hrx::ValueId(11), @@ -457,7 +616,7 @@ static void run_command_plan_metadata_checks() { REQUIRE(found_bundle->expert_table == ggml::hrx::ValueId(12)); REQUIRE(found_bundle->partition_table == ggml::hrx::ValueId(13)); ggml::hrx::CommandPlanMoeRoutingBundle conflicting_bundle = bundle; - conflicting_bundle.route_weights = ggml::hrx::ValueId(14); + conflicting_bundle.route_weights = ggml::hrx::ValueId(14); REQUIRE(!bundle_plan.append_moe_routing_bundle(conflicting_bundle, bundle_status)); REQUIRE(!bundle_status.success()); } @@ -472,6 +631,41 @@ static bool has_dispatch_registration(const std::vector & registrations, + const char * name, + ggml::hrx::DispatchMatchKind kind) { + for (const ggml::hrx::DispatchRegistration & registration : registrations) { + if (std::string(registration.name) == name && registration.kind == kind) { + return true; + } + } + return false; +} + +static void restore_environment_value(const char * name, bool had_value, const std::string & value) { + if (had_value) { + REQUIRE(setenv(name, value.c_str(), 1) == 0); + } else { + REQUIRE(unsetenv(name) == 0); + } +} + +static void add_binary_f32_exact_dispatch_params(ggml::hrx::Dispatch & dispatch, int64_t element_count) { + dispatch.kernel.integer_parameters.emplace("element_count", element_count); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", "1"); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", "1"); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride3", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src0_span", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.src1_span", std::to_string(element_count)); + dispatch.kernel.compile_parameters.emplace("ggml.binary_f32.op", "0"); +} + static bool match_test_single_dispatch(const ggml::hrx::DispatchMatchContext & context, ggml::hrx::DispatchMatch & match) { ggml::hrx::Dispatch dispatch; @@ -500,39 +694,182 @@ static bool match_test_wrong_root_dispatch(const ggml::hrx::DispatchMatchContext } static void run_dispatch_registry_checks() { + 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; + const std::string original_env_value = had_original_env ? original_env : ""; + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + const ggml::hrx::DispatchRegistry & registry = test_dispatch_registry(); REQUIRE(ggml::hrx::find_dispatch_registry({ "gfx1100" }) != nullptr); REQUIRE(ggml::hrx::find_dispatch_registry({ "gfx1151" }) != nullptr); REQUIRE(ggml::hrx::find_dispatch_registry({ "gfx0000" }) == nullptr); - REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_ADD), "common.add_f32")); - REQUIRE( - has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.matmul.dense_q4k_f16_wmma")); - REQUIRE( - has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.matmul.dense_q6k_f16_wmma")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_ADD), "common.binary_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_SUB), "common.binary_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL), "common.binary_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_DIV), "common.binary_f32")); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_swiglu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.packed_mul_mat_glu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_add.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_unary.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_unary.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.tiled_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.skinny_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "qwen.matmul.q6k_q8_1_x4")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.moe_router.projection_f32_four_row_wave32")); - REQUIRE( - has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), "qwen.rmsnorm_f32.mul_weight")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), "qwen.rmsnorm_f32_quantize_q8_1_x4")); + REQUIRE( + has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), "common.rmsnorm_binary_q8_1_x4")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), "common.rmsnorm_binary_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_RMS_NORM), "common.rmsnorm_f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), + "common.flash_attention_f32_f16_wmma")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), - "qwen.flash_attention_f32_f16_wmma")); + "common.flash_attention_decode_split_next_q8")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_RESHAPE), "qwen.attention_postprocess_f32_f16")); - REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), - "qwen.preamble.token_embedding_q4k")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32_next")); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows_scale.f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.gather_add_f32")); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope_set_rows.f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope_concat.f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope.f32", + ggml::hrx::DispatchMatchKind::SingleOp)); + REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_SET_ROWS), "common.set_rows")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_SOFT_MAX), "llm.moe_router.top8_f32")); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "common.mul_mat_id_swiglu.f32_f32_wmma", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "common.mul_mat_id_postops.f32_f32_wmma", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "common.mul_mat_id.f32_f32_wmma", ggml::hrx::DispatchMatchKind::SingleOp)); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), - "llm.routed_ffn.gate_up_swiglu_q4k_f16_wmma")); + "llm.routed_ffn.gate_up_swiglu_f16_wmma")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), "llm.routed_ffn.down_q4k_f16_wmma_grouped")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT_ID), "llm.routed_ffn.down_q6k_f16_wmma_grouped")); REQUIRE(registry.single_op_registrations().size() >= 5); + REQUIRE(setenv(kDisableQwenDispatchEnv, "0", 1) == 0); + const ggml::hrx::DispatchRegistry & false_env_registry = test_dispatch_registry(); + REQUIRE( + has_dispatch_registration(false_env_registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32")); + + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + const ggml::hrx::DispatchRegistry & generic_registry = test_dispatch_registry(); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_ADD), "common.binary_f32")); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RMS_NORM), + "common.rmsnorm_binary_f32")); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RMS_NORM), + "common.rmsnorm_binary_q8_1_x4")); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RMS_NORM), "common.rmsnorm_f32")); + REQUIRE( + has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_GET_ROWS), "common.gather_add_f32")); + REQUIRE( + has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32")); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_GET_ROWS), + "common.get_rows.f32_next")); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_ROPE), + "common.rope_set_rows.f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_ROPE), "common.rope.f32", + ggml::hrx::DispatchMatchKind::SingleOp)); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_SET_ROWS), "common.set_rows")); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RMS_NORM), + "qwen.rmsnorm_f32_quantize_q8_1_x4")); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_swiglu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.packed_mul_mat_glu.tiled_pair_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_add.skinny_f32_f32_decode", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_unary.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_unary.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.tiled_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "common.mul_mat_postops.skinny_f32_f32", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), + "qwen.matmul.q6k_q8_1_x4")); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), + "common.flash_attention_f32_f16_wmma")); + REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), + "common.flash_attention_decode_split_next_q8")); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RESHAPE), + "qwen.attention_postprocess_f32_f16")); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_GET_ROWS), + "qwen.preamble.token_embedding_q4k")); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_SOFT_MAX), + "llm.moe_router.top8_f32")); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "common.mul_mat_id_swiglu.f32_f32_wmma", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "common.mul_mat_id_postops.f32_f32_wmma", + ggml::hrx::DispatchMatchKind::Fused)); + REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "common.mul_mat_id.f32_f32_wmma", ggml::hrx::DispatchMatchKind::SingleOp)); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "llm.routed_ffn.gate_up_swiglu_f16_wmma")); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "llm.routed_ffn.down_q4k_f16_wmma_grouped")); + REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_MUL_MAT_ID), + "llm.routed_ffn.down_q6k_f16_wmma_grouped")); + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); + ggml::hrx::DispatchRegistryBuilder builder; builder.add({ "test.single_add", @@ -579,14 +916,16 @@ static void run_dispatch_registry_checks() { static void run_graph_import_checks() { ggml_init_params params = {}; - params.mem_size = 256 * 1024; + params.mem_size = 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * out = ggml_add(ctx, a, a); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out = ggml_add(ctx, a, b); REQUIRE(a != nullptr); + REQUIRE(b != nullptr); REQUIRE(out != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); @@ -599,21 +938,26 @@ static void run_graph_import_checks() { const ggml::hrx::GraphNode * node = &imported.graph.nodes().front(); REQUIRE(node->op == GGML_OP_ADD); REQUIRE(node->inputs.size() == 2); - REQUIRE(node->inputs[0] == node->inputs[1]); + REQUIRE(node->inputs[0] != node->inputs[1]); REQUIRE(ggml::hrx::DispatchScheduler::supports_node(imported.graph, node, test_dispatch_target())); const ggml::hrx::Value * a_value = imported.graph.values().find_tensor(a); + const ggml::hrx::Value * b_value = imported.graph.values().find_tensor(b); const ggml::hrx::Value * out_value = imported.graph.values().find_tensor(out); REQUIRE(a_value != nullptr); + REQUIRE(b_value != nullptr); REQUIRE(out_value != nullptr); REQUIRE(a_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(b_value->kind == ggml::hrx::ValueKind::External); REQUIRE(out_value->kind == ggml::hrx::ValueKind::External); REQUIRE(!a_value->buffer.has_value()); + REQUIRE(!b_value->buffer.has_value()); REQUIRE(!out_value->buffer.has_value()); const std::vector external_ids = imported.graph.values().external_value_ids(); - REQUIRE(external_ids.size() == 2); + REQUIRE(external_ids.size() == 3); REQUIRE(contains_value_id(external_ids, a_value->id)); + REQUIRE(contains_value_id(external_ids, b_value->id)); REQUIRE(contains_value_id(external_ids, out_value->id)); ggml::hrx::DispatchScheduler scheduler; @@ -632,6 +976,11 @@ static void run_graph_import_checks() { ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); REQUIRE(!partial_bindings.valid()); + REQUIRE(imported.graph.values().bind_buffer(b_value->id, { dummy_hrx_buffer(0x1800), 0, b_value->byte_count })); + ggml::hrx::CommandProgramBindings missing_output_binding = + ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); + REQUIRE(!missing_output_binding.valid()); + REQUIRE(imported.graph.values().bind_buffer(out_value->id, { dummy_hrx_buffer(0x2000), 0, 0 })); ggml::hrx::CommandProgramBindings empty_runtime_binding = ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); @@ -641,12 +990,17 @@ static void run_graph_import_checks() { ggml::hrx::CommandProgramBindings runtime_bindings = ggml::hrx::CommandProgramBindings::from_value_map(imported.graph.values()); REQUIRE(runtime_bindings.valid()); - REQUIRE(runtime_bindings.bindings().size() == 2); + REQUIRE(runtime_bindings.bindings().size() == 3); const ggml::hrx::CommandProgramBinding * a_binding = runtime_bindings.find(a_value->id); REQUIRE(a_binding != nullptr); REQUIRE(a_binding->buffer == dummy_hrx_buffer(0x1000)); REQUIRE(a_binding->offset == 0); REQUIRE(a_binding->length == a_value->byte_count); + const ggml::hrx::CommandProgramBinding * b_binding = runtime_bindings.find(b_value->id); + REQUIRE(b_binding != nullptr); + REQUIRE(b_binding->buffer == dummy_hrx_buffer(0x1800)); + REQUIRE(b_binding->offset == 0); + REQUIRE(b_binding->length == b_value->byte_count); REQUIRE(runtime_bindings.find(ggml::hrx::ValueId(123456)) == nullptr); const ggml::hrx::CommandProgramBindingsFingerprint runtime_fingerprint = @@ -689,10 +1043,10 @@ static void run_graph_import_checks() { REQUIRE(command.kind == ggml::hrx::CommandKind::Kernel); REQUIRE(command.kernel.kernel_id != ggml::hrx::kUncatalogedKernelId); REQUIRE(command.bindings.size() == 3); - REQUIRE(command.bindings[0].name == "a"); + REQUIRE(command.bindings[0].name == "lhs"); REQUIRE(command.bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); REQUIRE(command.bindings[0].access == ggml::hrx::ResourceAccess::Read); - REQUIRE(command.bindings[1].name == "b"); + REQUIRE(command.bindings[1].name == "rhs"); REQUIRE(command.bindings[1].origin == ggml::hrx::CommandBindingOrigin::GraphValue); REQUIRE(command.bindings[1].access == ggml::hrx::ResourceAccess::Read); REQUIRE(command.bindings[2].name == "output"); @@ -717,7 +1071,7 @@ static void run_graph_import_checks() { REQUIRE(ggml::hrx::resource_access_name(static_cast(255)) == "Unknown(255)"); const std::string binding_text = ggml::hrx::format_command_binding(command.bindings[0]); - REQUIRE(string_contains(binding_text, "binding a")); + REQUIRE(string_contains(binding_text, "binding lhs")); REQUIRE(string_contains(binding_text, "value=")); REQUIRE(string_contains(binding_text, "origin=GraphValue")); REQUIRE(string_contains(binding_text, "access=Read")); @@ -737,8 +1091,8 @@ static void run_graph_import_checks() { const std::string program_text = ggml::hrx::format_command_program(commands); REQUIRE(string_contains(program_text, "command_program commands=1")); REQUIRE(string_contains(program_text, "command 0")); - REQUIRE(string_contains(program_text, "binding a")); - REQUIRE(string_contains(program_text, "binding b")); + REQUIRE(string_contains(program_text, "binding lhs")); + REQUIRE(string_contains(program_text, "binding rhs")); REQUIRE(string_contains(program_text, "binding output")); ggml::hrx::ResolvedCommandProgram resolved = @@ -749,14 +1103,14 @@ static void run_graph_import_checks() { REQUIRE(resolved.commands.front().kind == command.kind); REQUIRE(resolved.commands.front().kernel.kernel_id == command.kernel.kernel_id); REQUIRE(resolved.commands.front().bindings.size() == 3); - REQUIRE(resolved.commands.front().bindings[0].binding.name == "a"); + REQUIRE(resolved.commands.front().bindings[0].binding.name == "lhs"); REQUIRE(resolved.commands.front().bindings[0].ref.buffer == dummy_hrx_buffer(0x1000)); REQUIRE(resolved.commands.front().bindings[0].ref.offset == 0); REQUIRE(resolved.commands.front().bindings[0].ref.length == a_value->byte_count); - REQUIRE(resolved.commands.front().bindings[1].binding.name == "b"); - REQUIRE(resolved.commands.front().bindings[1].ref.buffer == dummy_hrx_buffer(0x1000)); + REQUIRE(resolved.commands.front().bindings[1].binding.name == "rhs"); + REQUIRE(resolved.commands.front().bindings[1].ref.buffer == dummy_hrx_buffer(0x1800)); REQUIRE(resolved.commands.front().bindings[1].ref.offset == 0); - REQUIRE(resolved.commands.front().bindings[1].ref.length == a_value->byte_count); + REQUIRE(resolved.commands.front().bindings[1].ref.length == b_value->byte_count); REQUIRE(resolved.commands.front().bindings[2].binding.name == "output"); REQUIRE(resolved.commands.front().bindings[2].ref.buffer == dummy_hrx_buffer(0x2000)); REQUIRE(resolved.commands.front().bindings[2].ref.offset == 0); @@ -824,7 +1178,7 @@ static void run_graph_import_checks() { resolved = ggml::hrx::resolve_command_program_bindings(commands, null_bindings); REQUIRE(!resolved.valid()); REQUIRE(status_contains(resolved.status, "null buffer")); - REQUIRE(status_contains(resolved.status, "binding a")); + REQUIRE(status_contains(resolved.status, "binding lhs")); ggml::hrx::CommandProgram empty_resolve_binding = copy_command_program_shape(commands); empty_resolve_binding.commands.front().bindings[0].length = 0; @@ -839,7 +1193,7 @@ static void run_graph_import_checks() { resolved = ggml::hrx::resolve_command_program_bindings(out_of_range_binding, runtime_bindings); REQUIRE(!resolved.valid()); REQUIRE(status_contains(resolved.status, "outside runtime binding length")); - REQUIRE(status_contains(resolved.status, "binding a")); + REQUIRE(status_contains(resolved.status, "binding lhs")); ggml::hrx::CommandProgram unsupported_origin = copy_command_program_shape(commands); unsupported_origin.commands.front().bindings[0].origin = static_cast(255); @@ -870,7 +1224,7 @@ static void run_graph_import_checks() { empty_binding.commands.front().bindings[0].length = 0; verification = ggml::hrx::verify_command_program(empty_binding, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(!verification.valid()); - REQUIRE(status_contains(verification.status, "binding a")); + REQUIRE(status_contains(verification.status, "binding lhs")); REQUIRE(status_contains(verification.status, "range=[0, 0)")); ggml::hrx::CommandProgram invalid_value = copy_command_program_shape(commands); @@ -900,6 +1254,40 @@ static void run_graph_import_checks() { REQUIRE(!verification.valid()); REQUIRE(status_contains(verification.status, "access=Write")); + ggml::hrx::CommandProgram transient_flow = copy_command_program_shape(commands); + transient_flow.commands.resize(2); + transient_flow.commands[0].ordinal = 0; + transient_flow.commands[0].dependencies.clear(); + transient_flow.commands[1] = transient_flow.commands[0]; + transient_flow.commands[1].ordinal = 1; + transient_flow.commands[1].dependencies.push_back(0); + const ggml::hrx::ValueId transient_value(100000); + transient_flow.transients.allocations.push_back({ + transient_value, + command.bindings[2].length, + 256, + 0, + }); + transient_flow.transients.arena_size = command.bindings[2].length; + transient_flow.commands[0].bindings[2].origin = ggml::hrx::CommandBindingOrigin::Transient; + transient_flow.commands[0].bindings[2].value = transient_value; + transient_flow.commands[1].bindings[0].origin = ggml::hrx::CommandBindingOrigin::Transient; + transient_flow.commands[1].bindings[0].value = transient_value; + transient_flow.commands[1].bindings[0].length = command.bindings[2].length; + verification = ggml::hrx::verify_command_program(transient_flow, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(verification.valid()); + + ggml::hrx::CommandProgram read_before_write = copy_command_program_shape(transient_flow); + std::swap(read_before_write.commands[0], read_before_write.commands[1]); + read_before_write.commands[0].ordinal = 0; + read_before_write.commands[0].dependencies.clear(); + read_before_write.commands[1].ordinal = 1; + read_before_write.commands[1].dependencies.clear(); + read_before_write.commands[1].dependencies.push_back(0); + verification = ggml::hrx::verify_command_program(read_before_write, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(!verification.valid()); + REQUIRE(status_contains(verification.status, "reads transient value 100000 before write by command 1")); + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); const ggml::hrx::CommandProgramExecutionContext prepare_context = { nullptr, nullptr, "gfx1151", &corpus, nullptr, nullptr, nullptr, nullptr, @@ -928,236 +1316,497 @@ static void run_graph_import_checks() { ggml_free(ctx); } -static void run_graph_snapshot_diagnostics_checks() { +static void run_graph_import_mixed_backend_boundary_checks() { + ggml_init_params params = {}; + params.mem_size = 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, 8, 4); + ggml_tensor * cpu_sin = ggml_sin(ctx, input); + ggml_tensor * hrx_norm = ggml_rms_norm(ctx, cpu_sin, 1.0e-6f); + REQUIRE(input != nullptr); + REQUIRE(cpu_sin != nullptr); + REQUIRE(hrx_norm != nullptr); + + ggml_cgraph * cpu_to_hrx = ggml_new_graph(ctx); + REQUIRE(cpu_to_hrx != nullptr); + ggml_graph_add_node(cpu_to_hrx, hrx_norm); + + ggml::hrx::GraphImportResult imported_cpu_to_hrx = ggml::hrx::import_ggml_graph(*cpu_to_hrx); + REQUIRE(imported_cpu_to_hrx.valid()); + REQUIRE(imported_cpu_to_hrx.graph.nodes().size() == 1); + REQUIRE(imported_cpu_to_hrx.graph.nodes()[0].op == GGML_OP_RMS_NORM); + + const ggml::hrx::Value * cpu_sin_value = imported_cpu_to_hrx.graph.values().find_tensor(cpu_sin); + const ggml::hrx::Value * hrx_norm_value = imported_cpu_to_hrx.graph.values().find_tensor(hrx_norm); + REQUIRE(cpu_sin_value != nullptr); + REQUIRE(hrx_norm_value != nullptr); + REQUIRE(cpu_sin_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(hrx_norm_value->kind == ggml::hrx::ValueKind::External); + + ggml_tensor * hrx_norm_for_cpu = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * cpu_sin_after = ggml_sin(ctx, hrx_norm_for_cpu); + REQUIRE(hrx_norm_for_cpu != nullptr); + REQUIRE(cpu_sin_after != nullptr); + + ggml_cgraph * hrx_to_cpu = ggml_new_graph(ctx); + REQUIRE(hrx_to_cpu != nullptr); + ggml_graph_add_node(hrx_to_cpu, hrx_norm_for_cpu); + + ggml::hrx::GraphImportResult imported_hrx_to_cpu = ggml::hrx::import_ggml_graph(*hrx_to_cpu); + REQUIRE(imported_hrx_to_cpu.valid()); + REQUIRE(imported_hrx_to_cpu.graph.nodes().size() == 1); + REQUIRE(imported_hrx_to_cpu.graph.nodes()[0].op == GGML_OP_RMS_NORM); + + const ggml::hrx::Value * input_value = imported_hrx_to_cpu.graph.values().find_tensor(input); + const ggml::hrx::Value * output_value = imported_hrx_to_cpu.graph.values().find_tensor(hrx_norm_for_cpu); + REQUIRE(input_value != nullptr); + REQUIRE(output_value != nullptr); + REQUIRE(input_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(output_value->kind == ggml::hrx::ValueKind::External); + + ggml_tensor * cpu_sin_for_view = ggml_sin(ctx, input); + ggml_tensor * cpu_sin_view = ggml_reshape_2d(ctx, cpu_sin_for_view, 8, 4); + ggml_tensor * hrx_norm_from_view = ggml_rms_norm(ctx, cpu_sin_view, 1.0e-6f); + REQUIRE(cpu_sin_for_view != nullptr); + REQUIRE(cpu_sin_view != nullptr); + REQUIRE(hrx_norm_from_view != nullptr); + + ggml_cgraph * cpu_view_to_hrx = ggml_new_graph(ctx); + REQUIRE(cpu_view_to_hrx != nullptr); + ggml_graph_add_node(cpu_view_to_hrx, cpu_sin_view); + ggml_graph_add_node(cpu_view_to_hrx, hrx_norm_from_view); + + ggml::hrx::GraphImportResult imported_cpu_view_to_hrx = ggml::hrx::import_ggml_graph(*cpu_view_to_hrx); + REQUIRE(imported_cpu_view_to_hrx.valid()); + REQUIRE(imported_cpu_view_to_hrx.graph.nodes().size() == 2); + + const ggml::hrx::Value * cpu_sin_for_view_value = + imported_cpu_view_to_hrx.graph.values().find_tensor(cpu_sin_for_view); + const ggml::hrx::Value * cpu_sin_view_value = imported_cpu_view_to_hrx.graph.values().find_tensor(cpu_sin_view); + REQUIRE(cpu_sin_for_view_value != nullptr); + REQUIRE(cpu_sin_view_value != nullptr); + REQUIRE(cpu_sin_for_view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(cpu_sin_view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(imported_cpu_view_to_hrx.graph.values().same_storage(cpu_sin_for_view_value->id, cpu_sin_view_value->id)); + + ggml_tensor * hrx_norm_for_view = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * hrx_norm_view = ggml_reshape_2d(ctx, hrx_norm_for_view, 8, 4); + ggml_tensor * cpu_sin_view_out = ggml_sin(ctx, hrx_norm_view); + REQUIRE(hrx_norm_for_view != nullptr); + REQUIRE(hrx_norm_view != nullptr); + REQUIRE(cpu_sin_view_out != nullptr); + + ggml_cgraph * hrx_view_to_cpu = ggml_new_graph(ctx); + REQUIRE(hrx_view_to_cpu != nullptr); + ggml_build_forward_expand(hrx_view_to_cpu, hrx_norm_view); + + ggml::hrx::GraphImportResult imported_hrx_view_to_cpu = ggml::hrx::import_ggml_graph(*hrx_view_to_cpu); + REQUIRE(imported_hrx_view_to_cpu.valid()); + + const ggml::hrx::Value * hrx_norm_for_view_value = + imported_hrx_view_to_cpu.graph.values().find_tensor(hrx_norm_for_view); + const ggml::hrx::Value * hrx_norm_view_value = imported_hrx_view_to_cpu.graph.values().find_tensor(hrx_norm_view); + REQUIRE(hrx_norm_for_view_value != nullptr); + REQUIRE(hrx_norm_view_value != nullptr); + REQUIRE(hrx_norm_for_view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(hrx_norm_view_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(imported_hrx_view_to_cpu.graph.values().same_storage(hrx_norm_for_view_value->id, hrx_norm_view_value->id)); + + ggml_free(ctx); +} + +static void run_graph_view_external_use_checks() { ggml_init_params params = {}; params.mem_size = 256 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * out = ggml_add(ctx, a, b); + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * shared = ggml_add(ctx, a, b); + ggml_tensor * left = ggml_mul(ctx, shared, c); + ggml_tensor * right = ggml_sqr(ctx, shared); + ggml_tensor * output = ggml_add(ctx, left, right); REQUIRE(a != nullptr); REQUIRE(b != nullptr); - REQUIRE(out != nullptr); + REQUIRE(c != nullptr); + REQUIRE(shared != nullptr); + REQUIRE(left != nullptr); + REQUIRE(right != nullptr); + REQUIRE(output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, out); + ggml_build_forward_expand(graph, output); + REQUIRE(graph->n_nodes == 4); + REQUIRE(graph->nodes[0] == shared); + REQUIRE(graph->nodes[1] == left); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + ggml_cgraph graph_view = ggml_graph_view(graph, 0, 2); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(graph_view); REQUIRE(imported.valid()); - const std::string snapshot_json = ggml::hrx::serialize_graph_snapshot_json(imported.graph, "gfx1151", 42); - REQUIRE(string_contains(snapshot_json, "ggml-hrx-graph-snapshot-v1")); - REQUIRE(string_contains(snapshot_json, "ADD")); + const ggml::hrx::Value * shared_value = imported.graph.values().find_tensor(shared); + const ggml::hrx::Value * left_value = imported.graph.values().find_tensor(left); + REQUIRE(shared_value != nullptr); + REQUIRE(left_value != nullptr); + REQUIRE(shared_value->kind == ggml::hrx::ValueKind::External); + REQUIRE(left_value->kind == ggml::hrx::ValueKind::External); - ggml::hrx::GraphSnapshotLoadResult loaded = ggml::hrx::load_graph_snapshot_json(snapshot_json); - REQUIRE(loaded.valid()); - REQUIRE(loaded.uid == 42); - REQUIRE(loaded.target == "gfx1151"); - REQUIRE(loaded.graph.nodes().size() == imported.graph.nodes().size()); - REQUIRE(loaded.graph.values().size() == imported.graph.values().size()); + ggml_free(ctx); +} - ggml::hrx::DispatchScheduler scheduler; - ggml::hrx::DispatchScheduleDiagnostics diagnostics; - REQUIRE(scheduler.schedule_graph(loaded.graph, { loaded.target }, &diagnostics)); +static ggml::hrx::Dispatch schedule_single_dispatch_for_tensor(ggml_context * ctx, ggml_tensor * output) { + 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); + return scheduler.plan().dispatches.front(); +} - ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - loaded.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), loaded.target); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); +static bool can_schedule_tensor(ggml_context * ctx, ggml_tensor * output) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); - ggml_free(ctx); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); } -static void run_unmatched_graph_diagnostics_checks() { +static void require_compile_param(const ggml::hrx::Dispatch & dispatch, const char * name, const char * value) { + REQUIRE(dispatch.kernel.compile_parameters.at(name) == value); +} + +static void run_scale_f32_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 512 * 1024; + params.mem_size = 256 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 512, 512); - ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 512, 2); - ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 2); - REQUIRE(cache != nullptr); - REQUIRE(rows != nullptr); - REQUIRE(indices != nullptr); - ggml_tensor * out = ggml_set_rows(ctx, cache, rows, indices); - REQUIRE(out != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 18); + ggml_tensor * output = ggml_scale(ctx, input, 1.75f); + REQUIRE(input != nullptr); + REQUIRE(output != nullptr); + + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_scale_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == 69120); + require_compile_parameter(dispatch, "ggml.scale_f32.scale", expected_config_value(1.75f)); + require_compile_parameter(dispatch, "ggml.scale_f32.bias", expected_config_value(0.0f)); + + ggml_tensor * biased = ggml_scale_bias(ctx, input, -0.5f, 0.25f); + REQUIRE(biased != nullptr); + + const ggml::hrx::Dispatch bias_dispatch = schedule_single_dispatch_for_tensor(ctx, biased); + REQUIRE(kernel_name_for_id(bias_dispatch.kernel.kernel_id) == "loom_libs:ggml_scale_f32"); + REQUIRE(bias_dispatch.kernel.integer_parameters.at("element_count") == 69120); + require_compile_parameter(bias_dispatch, "ggml.scale_f32.scale", expected_config_value(-0.5f)); + require_compile_parameter(bias_dispatch, "ggml.scale_f32.bias", expected_config_value(0.25f)); + + ggml_tensor * inplace = ggml_scale_bias_inplace(ctx, input, 2.0f, -1.0f); + REQUIRE(inplace != nullptr); + + const ggml::hrx::Dispatch inplace_dispatch = schedule_single_dispatch_for_tensor(ctx, inplace); + REQUIRE(kernel_name_for_id(inplace_dispatch.kernel.kernel_id) == "loom_libs:ggml_scale_bias_f32"); + REQUIRE(inplace_dispatch.kernel.integer_parameters.at("element_count") == 69120); + require_compile_parameter(inplace_dispatch, "ggml.scale.scale", expected_config_value(2.0f)); + require_compile_parameter(inplace_dispatch, "ggml.scale.bias", expected_config_value(-1.0f)); + + ggml_free(ctx); +} + +static void require_scale_add_dispatch(ggml_context * ctx, + ggml_tensor * input, + ggml_tensor * residual, + int64_t expected_element_count, + int64_t expected_input_stride1, + int64_t expected_residual_stride1, + ggml::hrx::BinaryKind binary_kind = ggml::hrx::BinaryKind::Add, + bool scaled_lhs = true) { + ggml_tensor * scaled = ggml_scale_bias(ctx, input, 0.177800179f, 0.25f); + ggml_tensor * output = nullptr; + switch (binary_kind) { + case ggml::hrx::BinaryKind::Add: + output = ggml_add(ctx, scaled, residual); + break; + case ggml::hrx::BinaryKind::Sub: + output = scaled_lhs ? ggml_sub(ctx, scaled, residual) : ggml_sub(ctx, residual, scaled); + break; + case ggml::hrx::BinaryKind::Mul: + output = ggml_mul(ctx, scaled, residual); + break; + case ggml::hrx::BinaryKind::Div: + output = scaled_lhs ? ggml_div(ctx, scaled, residual) : ggml_div(ctx, residual, scaled); + break; + default: + REQUIRE(false); + } + REQUIRE(scaled != nullptr); + REQUIRE(output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, out); + ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 1); - REQUIRE(imported.graph.nodes().front().op == GGML_OP_SET_ROWS); + require_scheduled_command_program(graph, [&](const ggml::hrx::Graph &, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_scale_add_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == expected_element_count); + require_compile_parameter(dispatch, "ggml.scale_add_f32.scale", expected_config_value(0.177800179f)); + require_compile_parameter(dispatch, "ggml.scale_add_f32.bias", expected_config_value(0.25f)); + require_compile_parameter(dispatch, "ggml.scale_add_f32.binary_op", + std::to_string(ggml::hrx::binary_kind_config_value(binary_kind))); + require_compile_parameter(dispatch, "ggml.scale_add_f32.scaled_lhs", scaled_lhs ? "1" : "0"); + require_compile_parameter(dispatch, "ggml.scale_add_f32.input_stride1", std::to_string(expected_input_stride1)); + require_compile_parameter(dispatch, "ggml.scale_add_f32.residual_stride1", + std::to_string(expected_residual_stride1)); + REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings[0].name == "input"); + REQUIRE(commands.commands.front().bindings[1].name == "residual"); + REQUIRE(commands.commands.front().bindings[2].name == "output"); + }); +} - ggml::hrx::DispatchScheduler scheduler; - ggml::hrx::DispatchScheduleDiagnostics diagnostics; - REQUIRE(!scheduler.schedule_graph(imported.graph, test_dispatch_target(), &diagnostics)); - REQUIRE(diagnostics.unsupported_node != nullptr); - REQUIRE(diagnostics.unsupported_node_index == 0); - REQUIRE(diagnostics.unsupported_node->op == GGML_OP_SET_ROWS); - REQUIRE(diagnostics.match.attempts.empty()); - REQUIRE(status_contains(scheduler.plan().status, "unsupported HRX node 0: SET_ROWS")); +static void run_scale_add_f32_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 std::string diagnostics_text = - ggml::hrx::format_schedule_diagnostics_text(imported.graph, scheduler.plan(), diagnostics); - REQUIRE(string_contains(diagnostics_text, "unsupported_node=0:SET_ROWS")); - REQUIRE(string_contains(diagnostics_text, "matcher_attempts=0")); - const std::string diagnostics_json = - ggml::hrx::serialize_schedule_diagnostics_json(imported.graph, scheduler.plan(), diagnostics); - REQUIRE(string_contains(diagnostics_json, "SET_ROWS")); - REQUIRE(string_contains(diagnostics_json, "matcher_attempts")); + ggml_tensor * packed_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); + ggml_tensor * packed_residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); + REQUIRE(packed_input != nullptr); + REQUIRE(packed_residual != nullptr); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Sub, + false); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Mul, + true); + require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Div, + true); + + ggml_tensor * input_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 320, 3); + ggml_tensor * residual_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 384, 3); + ggml_tensor * strided_input = ggml_view_2d(ctx, input_storage, 257, 3, input_storage->nb[1], 0); + ggml_tensor * strided_residual = ggml_view_2d(ctx, residual_storage, 257, 3, residual_storage->nb[1], 0); + REQUIRE(strided_input != nullptr); + REQUIRE(strided_residual != nullptr); + require_scale_add_dispatch(ctx, strided_input, strided_residual, 771, 320, 384); + + ggml_tensor * scaled = ggml_scale(ctx, packed_input, 0.5f); + ggml_tensor * alias_output = ggml_add_inplace(ctx, scaled, packed_residual); + REQUIRE(scaled != nullptr); + REQUIRE(alias_output != nullptr); + REQUIRE(!can_schedule_tensor(ctx, alias_output)); + + ggml_tensor * alias_scaled = ggml_scale_inplace(ctx, packed_input, 0.5f); + ggml_tensor * output_from_alias = ggml_add(ctx, alias_scaled, packed_residual); + REQUIRE(alias_scaled != nullptr); + REQUIRE(output_from_alias != nullptr); + ggml_cgraph * alias_graph = ggml_new_graph(ctx); + REQUIRE(alias_graph != nullptr); + ggml_build_forward_expand(alias_graph, output_from_alias); + ggml::hrx::GraphImportResult alias_imported = ggml::hrx::import_ggml_graph(*alias_graph); + REQUIRE(alias_imported.valid()); + ggml::hrx::DispatchScheduler alias_scheduler; + REQUIRE(alias_scheduler.schedule_graph(alias_imported.graph, test_dispatch_target())); + REQUIRE(alias_scheduler.plan().dispatches.size() == 2); + REQUIRE(kernel_name_for_id(alias_scheduler.plan().dispatches[0].kernel.kernel_id) == + "loom_libs:ggml_scale_bias_f32"); + REQUIRE(kernel_name_for_id(alias_scheduler.plan().dispatches[1].kernel.kernel_id) == "loom_libs:ggml_binary_f32"); ggml_free(ctx); } -static void run_completion_counter_plan_checks() { - ggml::hrx::Graph graph; - ggml::hrx::CommandPlan plan; - const ggml::hrx::ValueId first_counter(1000); - const ggml::hrx::ValueId second_counter(1001); - plan.completion_counter_requests.push_back({ first_counter, "first_completion_counter", 1 }); - plan.completion_counter_requests.push_back({ second_counter, "second_completion_counters", 2 }); +static ggml_tensor * make_test_binary(ggml_context * ctx, + ggml_tensor * lhs, + ggml_tensor * rhs, + ggml::hrx::BinaryKind kind) { + switch (kind) { + case ggml::hrx::BinaryKind::Add: + return ggml_add(ctx, lhs, rhs); + case ggml::hrx::BinaryKind::Sub: + return ggml_sub(ctx, lhs, rhs); + case ggml::hrx::BinaryKind::Mul: + return ggml_mul(ctx, lhs, rhs); + case ggml::hrx::BinaryKind::Div: + return ggml_div(ctx, lhs, rhs); + default: + REQUIRE(false); + return nullptr; + } +} - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.empty()); - REQUIRE(commands.constant_initializations.empty()); - REQUIRE(commands.completion_counters.count == 3); - REQUIRE(commands.completion_counters.arena_offset == 0); - REQUIRE(commands.completion_counters.byte_count == 24); - REQUIRE(commands.transients.allocations.size() == 2); - REQUIRE(commands.transients.arena_size == 256); - REQUIRE(command_program_verifies(commands)); +static const char * expected_plain_rmsnorm_binary_kernel(int64_t hidden_size, int64_t token_count) { + return token_count < 6 && hidden_size >= 2048 && hidden_size <= 5120 ? + "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" : + "loom_libs:ggml_rmsnorm_binary_f32"; +} - const ggml::hrx::TransientAllocation * first_allocation = - ggml::hrx::find_transient_allocation(commands.transients, first_counter); - const ggml::hrx::TransientAllocation * second_allocation = - ggml::hrx::find_transient_allocation(commands.transients, second_counter); - REQUIRE(first_allocation != nullptr); - REQUIRE(second_allocation != nullptr); - REQUIRE(first_allocation->size == sizeof(int32_t)); - REQUIRE(first_allocation->alignment == 16); - REQUIRE(first_allocation->arena_offset == commands.completion_counters.arena_offset); - REQUIRE(second_allocation->size == 2 * sizeof(int32_t)); - REQUIRE(second_allocation->alignment == 16); - REQUIRE(second_allocation->arena_offset == 16); +static void run_rmsnorm_lowtoken_striped_dispatch_checks() { + const struct { + int64_t hidden_size; + int64_t token_count; + ggml::hrx::BinaryKind op; + const char * expected_kernel; + } cases[] = { + { 1024, 1, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32" }, + { 2048, 1, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, + { 3072, 2, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, + { 5120, 4, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, + { 5120, 6, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32" }, + { 2048, 1, ggml::hrx::BinaryKind::Add, "loom_libs:ggml_rmsnorm_binary_f32" }, + }; - ggml::hrx::CommandPlan bound_plan; - bound_plan.completion_counter_requests.push_back({ first_counter, "first_bound_completion_counter", 1 }); - bound_plan.completion_counter_requests.push_back({ second_counter, "second_bound_completion_counter", 1 }); - ggml::hrx::Dispatch first_dispatch; - first_dispatch.kernel = - ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref("qwen3_moe", "ggml_add_f32")); - first_dispatch.kernel.integer_parameters.emplace("element_count", 1); - first_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); - first_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); - first_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); - bound_plan.dispatches.push_back(std::move(first_dispatch)); - ggml::hrx::Dispatch second_dispatch; - second_dispatch.kernel = - ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref("qwen3_moe", "ggml_add_f32")); - second_dispatch.kernel.integer_parameters.emplace("element_count", 1); - second_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); - second_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); - second_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); - bound_plan.dispatches.push_back(std::move(second_dispatch)); + for (const auto & test : cases) { + ggml_init_params params = {}; + params.mem_size = static_cast(test.hidden_size * test.token_count * sizeof(float) * 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, test.hidden_size, test.token_count); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden_size); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * output = make_test_binary(ctx, rms, weight, test.op); + REQUIRE(output != nullptr); - const ggml::hrx::CommandProgram bound_commands = - ggml::hrx::build_command_program(graph, bound_plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(bound_commands.valid()); - REQUIRE(bound_commands.commands.size() == 2); - REQUIRE(bound_commands.completion_counters.count == 2); - REQUIRE(bound_commands.completion_counters.arena_offset == 0); - REQUIRE(bound_commands.completion_counters.byte_count == 20); - REQUIRE(command_program_verifies(bound_commands)); + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == test.expected_kernel); + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "ggml.rmsnorm_binary_f32.hidden_size", + std::to_string(test.hidden_size)); + require_compile_parameter(dispatch, "ggml.rmsnorm_binary_f32.op", + std::to_string(ggml::hrx::binary_kind_config_value(test.op))); + ggml_free(ctx); + } +} - const ggml::hrx::TransientArenaAllocationRef transient_arena = { - dummy_hrx_buffer(0x8000), - bound_commands.transients.arena_size, - 7, +static void run_rmsnorm_two_binary_dispatch_checks() { + const struct { + ggml::hrx::BinaryKind normalized_op; + ggml::hrx::BinaryKind output_op; + bool binary_lhs; + } cases[] = { + { ggml::hrx::BinaryKind::Add, ggml::hrx::BinaryKind::Mul, true }, + { ggml::hrx::BinaryKind::Sub, ggml::hrx::BinaryKind::Sub, true }, + { ggml::hrx::BinaryKind::Mul, ggml::hrx::BinaryKind::Div, true }, + { ggml::hrx::BinaryKind::Div, ggml::hrx::BinaryKind::Sub, false }, }; - ggml::hrx::PreparedCommandProgram prepared_shape; - for (const ggml::hrx::Command & prepared_source : bound_commands.commands) { - ggml::hrx::PreparedCommand prepared_command; - prepared_command.ordinal = prepared_source.ordinal; - prepared_command.kind = prepared_source.kind; - prepared_command.kernel.specialization = prepared_source.kernel; - for (const ggml::hrx::CommandBinding & binding : prepared_source.bindings) { - prepared_command.kernel.bindings.push_back({ - binding, { dummy_hrx_buffer(0x4000), 123, binding.length } - }); - } - prepared_shape.commands.push_back(std::move(prepared_command)); - } - prepared_shape.bound_transient_arena_allocation_id = 1; - REQUIRE(ggml::hrx::bind_prepared_command_program_transients(bound_commands, transient_arena, prepared_shape)); - REQUIRE(prepared_shape.commands[0].kernel.bindings[0].ref.buffer == dummy_hrx_buffer(0x8000)); - REQUIRE(prepared_shape.commands[0].kernel.bindings[0].ref.offset == 0); - REQUIRE(prepared_shape.commands[0].kernel.bindings[1].ref.buffer == dummy_hrx_buffer(0x8000)); - REQUIRE(prepared_shape.commands[0].kernel.bindings[1].ref.offset == 16); - REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.offset == 16); - REQUIRE(prepared_shape.commands[1].kernel.bindings[1].ref.offset == 0); + 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); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * scaled = make_test_binary(ctx, rms, weight, test.normalized_op); + ggml_tensor * output = test.binary_lhs ? make_test_binary(ctx, scaled, residual, test.output_op) : + make_test_binary(ctx, residual, scaled, test.output_op); + REQUIRE(output != nullptr); + + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_mul_add_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.normalized_op", + std::to_string(ggml::hrx::binary_kind_config_value(test.normalized_op))); + require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.output_op", + std::to_string(ggml::hrx::binary_kind_config_value(test.output_op))); + require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.binary_lhs", test.binary_lhs ? "1" : "0"); + ggml_free(ctx); + } } -static void run_graph_index_checks() { +static void run_cont_f32_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 256 * 1024; + params.mem_size = 512 * 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, 256, 1); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); - REQUIRE(input != nullptr); - REQUIRE(weight != nullptr); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); - REQUIRE(rms != nullptr); - ggml_tensor * out = ggml_mul(ctx, rms, weight); - REQUIRE(out != nullptr); + constexpr int64_t width = 64; + constexpr int64_t row_count = 40; + constexpr int64_t token_count = 23; + constexpr int64_t token_stride = 4096; + constexpr size_t view_offset = 8 * sizeof(float); + const size_t storage_elements = + view_offset / sizeof(float) + static_cast((token_count - 1) * token_stride + width * row_count); + + ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * view = ggml_view_3d(ctx, storage, width, row_count, token_count, width * sizeof(float), + token_stride * sizeof(float), view_offset); + ggml_tensor * output = ggml_cont(ctx, view); + REQUIRE(storage != nullptr); + REQUIRE(view != nullptr); + REQUIRE(output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, out); + ggml_build_forward_expand(graph, output); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); - REQUIRE(imported.graph.has_index()); REQUIRE(imported.graph.nodes().size() == 2); - const ggml::hrx::GraphNode * rms_node = &imported.graph.nodes()[0]; - const ggml::hrx::GraphNode * mul_node = &imported.graph.nodes()[1]; - REQUIRE(rms_node->op == GGML_OP_RMS_NORM); - REQUIRE(mul_node->op == GGML_OP_MUL); - const ggml::hrx::RmsNormParams * rms_params = ggml::hrx::op_params_as(rms_node->params); - REQUIRE(rms_params != nullptr); - REQUIRE(rms_params->eps == 0.000001f); - REQUIRE(imported.graph.index().producer(rms_node->output) == rms_node); - REQUIRE(imported.graph.index().producer(mul_node->output) == mul_node); - REQUIRE(imported.graph.index().has_single_consumer(rms_node->output)); - const std::vector & consumers = imported.graph.index().consumers(rms_node->output); - REQUIRE(consumers.size() == 1); - REQUIRE(consumers.front() == mul_node); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_VIEW); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_CONT); + + const ggml::hrx::Value * view_value = imported.graph.values().find(imported.graph.nodes()[1].inputs[0]); + REQUIRE(view_value != nullptr); + const ggml::hrx::Value * storage_value = imported.graph.values().find(view_value->storage_root); + REQUIRE(storage_value != nullptr); + REQUIRE(view_value->storage_offset == view_offset); ggml::hrx::DispatchScheduler scheduler; REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_strided_source_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == width * row_count * token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("source_span") == + static_cast((token_count - 1) * token_stride + width * row_count)); + require_compile_parameter(dispatch, "ggml.copy_strided_source_f32.ne0", std::to_string(width)); + require_compile_parameter(dispatch, "ggml.copy_strided_source_f32.ne1", std::to_string(row_count)); + require_compile_parameter(dispatch, "ggml.copy_strided_source_f32.ne2", std::to_string(token_count)); + require_compile_parameter(dispatch, "ggml.copy_strided_source_f32.stride0", "1"); + require_compile_parameter(dispatch, "ggml.copy_strided_source_f32.stride1", std::to_string(width)); + require_compile_parameter(dispatch, "ggml.copy_strided_source_f32.stride2", std::to_string(token_stride)); + REQUIRE(dispatch.bindings.size() == 2); + REQUIRE(dispatch.bindings[0].value == storage_value->id); + REQUIRE(dispatch.bindings[0].offset == view_offset); + REQUIRE(dispatch.bindings[0].length == + static_cast((token_count - 1) * token_stride + width * row_count) * sizeof(float)); + + 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)); ggml_free(ctx); } -static void run_graph_traversal_checks() { +static void run_binary_f32_broadcast_dispatch_checks() { { ggml_init_params params = {}; params.mem_size = 256 * 1024; @@ -1165,34 +1814,50 @@ static void run_graph_traversal_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * out0 = ggml_add(ctx, a, b); - ggml_tensor * out1 = ggml_add(ctx, c, d); - REQUIRE(a != nullptr); - REQUIRE(b != nullptr); - REQUIRE(c != nullptr); - REQUIRE(d != nullptr); - REQUIRE(out0 != nullptr); - REQUIRE(out1 != nullptr); + ggml_tensor * lhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 16, 3); + ggml_tensor * rhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 16, 3); + ggml_tensor * output = ggml_add(ctx, lhs, rhs); + REQUIRE(lhs != nullptr); + REQUIRE(rhs != nullptr); + REQUIRE(output != nullptr); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, out0); - ggml_build_forward_expand(graph, out1); + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_binary_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == 48); + require_compile_param(dispatch, "ggml.binary_f32.ne0", "16"); + require_compile_param(dispatch, "ggml.binary_f32.src0_stride1", "16"); + require_compile_param(dispatch, "ggml.binary_f32.src1_stride1", "16"); + REQUIRE(dispatch.kernel.integer_parameters.count("src0_element_count") == 0); + require_compile_param(dispatch, "ggml.binary_f32.op", "0"); + REQUIRE(dispatch.kernel.compile_parameters.count("ggml.binary_bc_f32.src0_broadcast_dim0") == 0); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 2); - REQUIRE(imported.graph.nodes()[0].output == imported.graph.values().find_tensor(out0)->id); - REQUIRE(imported.graph.nodes()[1].output == imported.graph.values().find_tensor(out1)->id); + ggml_free(ctx); + } - const std::vector order = traversal_indices(imported.graph); - REQUIRE(order.size() == 2); - REQUIRE(order[0] == 0); - REQUIRE(order[1] == 1); + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * source = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2048, 3, 2); + REQUIRE(source != nullptr); + ggml_tensor * lhs = ggml_view_2d(ctx, source, 2048, 2, source->nb[2], 0); + ggml_tensor * rhs = ggml_view_2d(ctx, source, 2048, 2, source->nb[2], 2 * source->nb[1]); + REQUIRE(lhs != nullptr); + REQUIRE(rhs != nullptr); + ggml_tensor * output = ggml_mul(ctx, lhs, rhs); + REQUIRE(output != nullptr); + + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_binary_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == 4096); + require_compile_param(dispatch, "ggml.binary_f32.src0_stride1", "6144"); + require_compile_param(dispatch, "ggml.binary_f32.src1_stride1", "6144"); + require_compile_param(dispatch, "ggml.binary_f32.src0_span", "8192"); + require_compile_param(dispatch, "ggml.binary_f32.src1_span", "8192"); + require_compile_param(dispatch, "ggml.binary_f32.op", "2"); ggml_free(ctx); } @@ -1204,34 +1869,54 @@ static void run_graph_traversal_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * add_out = ggml_add(ctx, a, b); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 3); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 2); - ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); - REQUIRE(a != nullptr); - REQUIRE(b != nullptr); - REQUIRE(add_out != nullptr); - REQUIRE(weight != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 16, 3); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16); + ggml_tensor * output = ggml_mul(ctx, input, weight); REQUIRE(input != nullptr); - REQUIRE(matmul != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(output != nullptr); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, add_out); - ggml_build_forward_expand(graph, matmul); + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_binary_bc_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == 48); + REQUIRE(dispatch.kernel.integer_parameters.at("ne0") == 16); + REQUIRE(dispatch.kernel.integer_parameters.at("ne1") == 3); + REQUIRE(dispatch.kernel.integer_parameters.at("src0_element_count") == 48); + REQUIRE(dispatch.kernel.integer_parameters.at("src1_element_count") == 16); + require_compile_param(dispatch, "ggml.binary_bc_f32.op", "2"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src0_broadcast_dim0", "0"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src0_broadcast_dim1", "0"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim0", "0"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim1", "1"); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 2); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); - REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); + ggml_free(ctx); + } - const std::vector order = traversal_indices(imported.graph); - REQUIRE(order.size() == 2); - REQUIRE(order[0] == 1); - REQUIRE(order[1] == 0); + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * numerator = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 4); + ggml_tensor * denominator = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 4); + ggml_tensor * output = ggml_div(ctx, numerator, denominator); + REQUIRE(numerator != nullptr); + REQUIRE(denominator != nullptr); + REQUIRE(output != nullptr); + + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_binary_bc_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == 32); + REQUIRE(dispatch.kernel.integer_parameters.at("ne0") == 8); + REQUIRE(dispatch.kernel.integer_parameters.at("ne1") == 4); + REQUIRE(dispatch.kernel.integer_parameters.at("src0_element_count") == 32); + REQUIRE(dispatch.kernel.integer_parameters.at("src1_element_count") == 4); + require_compile_param(dispatch, "ggml.binary_bc_f32.op", "3"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src0_broadcast_dim0", "0"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim0", "1"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim1", "0"); ggml_free(ctx); } @@ -1243,29 +1928,25 @@ static void run_graph_traversal_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 1); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); - ggml_tensor * out = ggml_mul(ctx, rms, weight); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 5, 6, 7); + ggml_tensor * weight = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 6, 7); + ggml_tensor * output = ggml_mul(ctx, input, weight); REQUIRE(input != nullptr); REQUIRE(weight != nullptr); - REQUIRE(rms != nullptr); - REQUIRE(out != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, out); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 2); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_RMS_NORM); - REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL); + REQUIRE(output != nullptr); - const std::vector order = traversal_indices(imported.graph); - REQUIRE(order.size() == 2); - REQUIRE(order[0] == 0); - REQUIRE(order[1] == 1); + const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_binary_bc_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == 210); + REQUIRE(dispatch.kernel.integer_parameters.at("ne0") == 5); + REQUIRE(dispatch.kernel.integer_parameters.at("ne1") == 6); + REQUIRE(dispatch.kernel.integer_parameters.at("ne2") == 7); + REQUIRE(dispatch.kernel.integer_parameters.at("src0_element_count") == 210); + REQUIRE(dispatch.kernel.integer_parameters.at("src1_element_count") == 42); + require_compile_param(dispatch, "ggml.binary_bc_f32.op", "2"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim0", "1"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim1", "0"); + require_compile_param(dispatch, "ggml.binary_bc_f32.src1_broadcast_dim2", "0"); ggml_free(ctx); } @@ -1277,1018 +1958,6735 @@ static void run_graph_traversal_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * left = ggml_add(ctx, a, b); - ggml_tensor * right = ggml_mul(ctx, a, c); - ggml_tensor * join = ggml_add(ctx, left, right); - REQUIRE(a != nullptr); - REQUIRE(b != nullptr); - REQUIRE(c != nullptr); - REQUIRE(left != nullptr); - REQUIRE(right != nullptr); - REQUIRE(join != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 4); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2, 4); + ggml_tensor * output = ggml_mul(ctx, input, weight); + REQUIRE(input != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(output != nullptr); + REQUIRE(!can_schedule_tensor(ctx, output)); + + ggml_free(ctx); + } +} + +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); - ggml_build_forward_expand(graph, join); + 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()); - REQUIRE(imported.graph.nodes().size() == 3); - - const size_t left_index = producer_index_for_tensor(imported.graph, left); - const size_t right_index = producer_index_for_tensor(imported.graph, right); - const size_t join_index = producer_index_for_tensor(imported.graph, join); - - const std::vector order = traversal_indices(imported.graph); - REQUIRE(order.size() == 3); - REQUIRE(find_position(order, join_index) > find_position(order, left_index)); - REQUIRE(find_position(order, join_index) > find_position(order, right_index)); + 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 schedule_single_matmul_command(ggml_context * ctx, - ggml_tensor * output, - const char * expected_kernel_name, - int64_t expected_token_count, - int64_t expected_input_size, - int64_t expected_output_size) { +static void run_graph_snapshot_diagnostics_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out = ggml_add(ctx, a, b); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(out != nullptr); + ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); + ggml_build_forward_expand(graph, out); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 1); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); + const std::string snapshot_json = ggml::hrx::serialize_graph_snapshot_json(imported.graph, "gfx1151", 42); + REQUIRE(string_contains(snapshot_json, "ggml-hrx-graph-snapshot-v1")); + REQUIRE(string_contains(snapshot_json, "ADD")); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == expected_kernel_name); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); - REQUIRE(dispatch.bindings.size() == 3); - require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); - if (string_contains(kernel_name, "dense_linear")) { - require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.input_size", - std::to_string(expected_input_size)); - require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_size", - std::to_string(expected_output_size)); - require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_accumulation", "0"); - } else { - require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", std::to_string(expected_input_size)); - require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expected_output_size)); - } + ggml::hrx::GraphSnapshotLoadResult loaded = ggml::hrx::load_graph_snapshot_json(snapshot_json); + REQUIRE(loaded.valid()); + REQUIRE(loaded.uid == 42); + REQUIRE(loaded.target == "gfx1151"); + REQUIRE(loaded.graph.nodes().size() == imported.graph.nodes().size()); + REQUIRE(loaded.graph.values().size() == imported.graph.values().size()); - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(loaded.graph, { loaded.target }, &diagnostics)); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + loaded.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), loaded.target); REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.front().bindings.size() == 3); - REQUIRE(commands.commands.front().bindings[0].name == "input"); - REQUIRE(commands.commands.front().bindings[1].name == "weight"); - REQUIRE(commands.commands.front().bindings[2].name == "output"); + + ggml_free(ctx); } -static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) { +static void run_unmatched_graph_diagnostics_checks() { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 512, 512); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 512, 2); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2); + REQUIRE(cache != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * out = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(out != nullptr); + ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); + ggml_build_forward_expand(graph, out); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); - return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes().front().op == GGML_OP_SET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(!scheduler.schedule_graph(imported.graph, test_dispatch_target(), &diagnostics)); + REQUIRE(diagnostics.unsupported_node != nullptr); + REQUIRE(diagnostics.unsupported_node_index == 0); + REQUIRE(diagnostics.unsupported_node->op == GGML_OP_SET_ROWS); + REQUIRE(!diagnostics.match.attempts.empty()); + REQUIRE(status_contains(scheduler.plan().status, "unsupported HRX node 0: SET_ROWS")); + + const std::string diagnostics_text = + ggml::hrx::format_schedule_diagnostics_text(imported.graph, scheduler.plan(), diagnostics); + REQUIRE(string_contains(diagnostics_text, "unsupported_node=0:SET_ROWS")); + REQUIRE(string_contains(diagnostics_text, "matcher_attempts=")); + const std::string diagnostics_json = + ggml::hrx::serialize_schedule_diagnostics_json(imported.graph, scheduler.plan(), diagnostics); + REQUIRE(string_contains(diagnostics_json, "SET_ROWS")); + REQUIRE(string_contains(diagnostics_json, "matcher_attempts")); + + ggml_free(ctx); } -static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { +static void run_completion_counter_plan_checks() { + ggml::hrx::Graph graph; + ggml::hrx::CommandPlan plan; + const ggml::hrx::ValueId first_counter(1000); + const ggml::hrx::ValueId second_counter(1001); + plan.completion_counter_requests.push_back({ first_counter, "first_completion_counter", 1 }); + plan.completion_counter_requests.push_back({ second_counter, "second_completion_counters", 2 }); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.empty()); + REQUIRE(commands.constant_initializations.empty()); + REQUIRE(commands.completion_counters.count == 3); + REQUIRE(commands.completion_counters.arena_offset == 0); + REQUIRE(commands.completion_counters.byte_count == 24); + REQUIRE(commands.transients.allocations.size() == 2); + REQUIRE(commands.transients.arena_size == 256); + REQUIRE(command_program_verifies(commands)); + + const ggml::hrx::TransientAllocation * first_allocation = + ggml::hrx::find_transient_allocation(commands.transients, first_counter); + const ggml::hrx::TransientAllocation * second_allocation = + ggml::hrx::find_transient_allocation(commands.transients, second_counter); + REQUIRE(first_allocation != nullptr); + REQUIRE(second_allocation != nullptr); + REQUIRE(first_allocation->size == sizeof(int32_t)); + REQUIRE(first_allocation->alignment == 16); + REQUIRE(first_allocation->arena_offset == commands.completion_counters.arena_offset); + REQUIRE(second_allocation->size == 2 * sizeof(int32_t)); + REQUIRE(second_allocation->alignment == 16); + REQUIRE(second_allocation->arena_offset == 16); + + ggml::hrx::CommandPlan bound_plan; + bound_plan.completion_counter_requests.push_back({ first_counter, "first_bound_completion_counter", 1 }); + bound_plan.completion_counter_requests.push_back({ second_counter, "second_bound_completion_counter", 1 }); + ggml::hrx::Dispatch first_dispatch; + first_dispatch.kernel = + ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref("loom_libs", "ggml_binary_f32")); + add_binary_f32_exact_dispatch_params(first_dispatch, 1); + first_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); + first_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); + first_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); + bound_plan.dispatches.push_back(std::move(first_dispatch)); + ggml::hrx::Dispatch second_dispatch; + second_dispatch.kernel = + ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref("loom_libs", "ggml_binary_f32")); + add_binary_f32_exact_dispatch_params(second_dispatch, 1); + second_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); + second_dispatch.bindings.push_back({ first_counter, 0, sizeof(int32_t) }); + second_dispatch.bindings.push_back({ second_counter, 0, sizeof(int32_t) }); + bound_plan.dispatches.push_back(std::move(second_dispatch)); + + const ggml::hrx::CommandProgram bound_commands = + ggml::hrx::build_command_program(graph, bound_plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(bound_commands.valid()); + REQUIRE(bound_commands.commands.size() == 2); + REQUIRE(bound_commands.completion_counters.count == 2); + REQUIRE(bound_commands.completion_counters.arena_offset == 0); + REQUIRE(bound_commands.completion_counters.byte_count == 20); + REQUIRE(command_program_verifies(bound_commands)); + + const ggml::hrx::TransientArenaAllocationRef transient_arena = { + dummy_hrx_buffer(0x8000), + bound_commands.transients.arena_size, + 7, + }; + ggml::hrx::PreparedCommandProgram prepared_shape; + for (const ggml::hrx::Command & prepared_source : bound_commands.commands) { + ggml::hrx::PreparedCommand prepared_command; + prepared_command.ordinal = prepared_source.ordinal; + prepared_command.kind = prepared_source.kind; + prepared_command.kernel.specialization = prepared_source.kernel; + for (const ggml::hrx::CommandBinding & binding : prepared_source.bindings) { + prepared_command.kernel.bindings.push_back({ + binding, { dummy_hrx_buffer(0x4000), 123, binding.length } + }); + } + prepared_shape.commands.push_back(std::move(prepared_command)); + } + prepared_shape.bound_transient_arena_allocation_id = 1; + + REQUIRE(ggml::hrx::bind_prepared_command_program_transients(bound_commands, transient_arena, prepared_shape)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[0].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[0].ref.offset == 0); + REQUIRE(prepared_shape.commands[0].kernel.bindings[1].ref.buffer == dummy_hrx_buffer(0x8000)); + REQUIRE(prepared_shape.commands[0].kernel.bindings[1].ref.offset == 16); + REQUIRE(prepared_shape.commands[1].kernel.bindings[0].ref.offset == 16); + REQUIRE(prepared_shape.commands[1].kernel.bindings[1].ref.offset == 0); +} + +static void run_graph_index_checks() { ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; + params.mem_size = 256 * 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, token_count); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); - ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); REQUIRE(input != nullptr); - REQUIRE(norm_weight != nullptr); - REQUIRE(vocab_weight != nullptr); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); - ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); - ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, normalized); + REQUIRE(weight != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); REQUIRE(rms != nullptr); - REQUIRE(normalized != nullptr); - REQUIRE(logits != nullptr); + ggml_tensor * out = ggml_mul(ctx, rms, weight); + REQUIRE(out != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, logits); + ggml_build_forward_expand(graph, out); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 3); + REQUIRE(imported.graph.has_index()); + REQUIRE(imported.graph.nodes().size() == 2); + const ggml::hrx::GraphNode * rms_node = &imported.graph.nodes()[0]; + const ggml::hrx::GraphNode * mul_node = &imported.graph.nodes()[1]; + REQUIRE(rms_node->op == GGML_OP_RMS_NORM); + REQUIRE(mul_node->op == GGML_OP_MUL); + const ggml::hrx::RmsNormParams * rms_params = ggml::hrx::op_params_as(rms_node->params); + REQUIRE(rms_params != nullptr); + REQUIRE(rms_params->eps == 0.000001f); + REQUIRE(imported.graph.index().producer(rms_node->output) == rms_node); + REQUIRE(imported.graph.index().producer(mul_node->output) == mul_node); + REQUIRE(imported.graph.index().has_single_consumer(rms_node->output)); + const std::vector & consumers = imported.graph.index().consumers(rms_node->output); + REQUIRE(consumers.size() == 1); + REQUIRE(consumers.front() == mul_node); 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); - - const ggml::hrx::Dispatch & rms_dispatch = scheduler.plan().dispatches[0]; - REQUIRE(kernel_name_for_id(rms_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); - REQUIRE(rms_dispatch.kernel.integer_parameters.at("token_count") == token_count); - require_compile_parameter(rms_dispatch, "qwen3_moe.model.hidden_size", "2048"); - require_compile_parameter(rms_dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); - REQUIRE(rms_dispatch.bindings.size() == 4); - - const ggml::hrx::CommandPlanTransient & q8_transient = scheduler.plan().transients.front(); - REQUIRE(q8_transient.size == qwen_q8_1_x4_size(token_count, 2048)); - REQUIRE(rms_dispatch.bindings[3].value == q8_transient.value); - REQUIRE(rms_dispatch.bindings[3].length == q8_transient.size); + REQUIRE(scheduler.plan().dispatches.size() == 1); - const ggml::hrx::Dispatch & vocab_dispatch = scheduler.plan().dispatches[1]; - REQUIRE(kernel_name_for_id(vocab_dispatch.kernel.kernel_id) == "qwen3_moe:ggml_linear_q6k_q8_1_x4"); - REQUIRE(vocab_dispatch.kernel.integer_parameters.at("token_count") == token_count); - REQUIRE(vocab_dispatch.kernel.integer_parameters.at("input_size") == 2048); - REQUIRE(vocab_dispatch.kernel.integer_parameters.at("output_size") == 151936); - require_compile_parameter(vocab_dispatch, "ggml.linear_q6k_q8_1_x4.token_capacity", std::to_string(token_count)); - require_compile_parameter(vocab_dispatch, "ggml.linear_q6k_q8_1_x4.output_capacity", "151936"); - REQUIRE(vocab_dispatch.bindings.size() == 3); - REQUIRE(vocab_dispatch.bindings[0].value == q8_transient.value); - REQUIRE(vocab_dispatch.bindings[0].length == q8_transient.size); + ggml_free(ctx); - 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() == 2); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands[0].bindings.size() == 4); - REQUIRE(commands.commands[0].bindings[3].name == "q8_output"); - REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); - REQUIRE(commands.commands[1].bindings.size() == 3); - REQUIRE(commands.commands[1].bindings[0].name == "q8_input"); - REQUIRE(commands.commands[1].bindings[0].origin == ggml::hrx::CommandBindingOrigin::Transient); - REQUIRE(commands.commands[1].bindings[0].value == q8_transient.value); + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ctx = ggml_init(params); + REQUIRE(ctx != nullptr); - ggml_free(ctx); -} + input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 4); + REQUIRE(input != nullptr); + rms = ggml_rms_norm(ctx, input, 0.00001f); + REQUIRE(rms != nullptr); -static bool graph_is_supported(ggml_context * ctx, ggml_tensor * output) { - ggml_cgraph * graph = ggml_new_graph(ctx); + graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); + ggml_build_forward_expand(graph, rms); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); - return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); -} + REQUIRE(imported.graph.nodes().size() == 1); -struct ManualQwenRouterTop8Graph { - ggml::hrx::Graph graph; - ggml::hrx::ValueId route_ids; -}; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == 4); + require_compile_parameter(dispatch, "ggml.rmsnorm_f32.hidden_size", "256"); + require_compile_parameter(dispatch, "ggml.rmsnorm_f32.input_stride", "256"); + REQUIRE(dispatch.kernel.compile_parameters.count("ggml.rmsnorm_f32.rms_epsilon") == 1); + REQUIRE(dispatch.bindings.size() == 2); -static ggml::hrx::ValueId add_manual_graph_tensor(ggml::hrx::Graph & graph, - ggml_tensor * tensor, - ggml::hrx::ValueKind kind) { + { + constexpr int64_t view_hidden_size = 256; + constexpr int64_t view_tokens = 23; + constexpr int64_t view_stride = 512; + constexpr size_t view_offset = 4 * sizeof(float); + ggml_tensor * storage = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, view_offset / sizeof(float) + (view_tokens - 1) * view_stride + view_hidden_size); + ggml_tensor * view = + ggml_view_2d(ctx, storage, view_hidden_size, view_tokens, view_stride * sizeof(float), view_offset); + ggml_tensor * norm = ggml_rms_norm(ctx, view, 0.00001f); + REQUIRE(storage != nullptr); + REQUIRE(view != nullptr); + REQUIRE(norm != nullptr); + + ggml_cgraph * view_graph = ggml_new_graph(ctx); + REQUIRE(view_graph != nullptr); + ggml_build_forward_expand(view_graph, norm); + + ggml::hrx::GraphImportResult view_imported = ggml::hrx::import_ggml_graph(*view_graph); + REQUIRE(view_imported.valid()); + REQUIRE(view_imported.graph.nodes().size() == 2); + REQUIRE(view_imported.graph.nodes()[0].op == GGML_OP_VIEW); + REQUIRE(view_imported.graph.nodes()[1].op == GGML_OP_RMS_NORM); + + const ggml::hrx::Value * view_value = + view_imported.graph.values().find(view_imported.graph.nodes()[1].inputs[0]); + REQUIRE(view_value != nullptr); + const ggml::hrx::Value * storage_value = view_imported.graph.values().find(view_value->storage_root); + REQUIRE(storage_value != nullptr); + REQUIRE(view_value->storage_offset == view_offset); + + REQUIRE(scheduler.schedule_graph(view_imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & view_dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(view_dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_f32"); + REQUIRE(view_dispatch.kernel.integer_parameters.at("token_count") == view_tokens); + require_compile_parameter(view_dispatch, "ggml.rmsnorm_f32.hidden_size", std::to_string(view_hidden_size)); + require_compile_parameter(view_dispatch, "ggml.rmsnorm_f32.input_stride", std::to_string(view_stride)); + REQUIRE(view_dispatch.bindings.size() == 2); + REQUIRE(view_dispatch.bindings[0].value == storage_value->id); + REQUIRE(view_dispatch.bindings[0].offset == view_offset); + REQUIRE(view_dispatch.bindings[0].length == static_cast(view_tokens - 1) * view->nb[1] + + static_cast(view_hidden_size) * sizeof(float)); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + view_imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + + { + constexpr int64_t hidden_size = 256; + constexpr int64_t token_count = 64; + constexpr int64_t input_stride = 288; + const size_t storage_elements = static_cast(token_count - 1) * input_stride + hidden_size; + ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * view = ggml_view_2d(ctx, storage, hidden_size, token_count, input_stride * sizeof(float), 0); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * norm = ggml_rms_norm(ctx, view, 0.00001f); + ggml_tensor * scaled = ggml_mul(ctx, norm, weight); + REQUIRE(storage != nullptr); + REQUIRE(view != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(norm != nullptr); + REQUIRE(scaled != nullptr); + + ggml_cgraph * strided_graph = ggml_new_graph(ctx); + REQUIRE(strided_graph != nullptr); + ggml_build_forward_expand(strided_graph, scaled); + ggml::hrx::GraphImportResult strided_imported = ggml::hrx::import_ggml_graph(*strided_graph); + REQUIRE(strided_imported.valid()); + REQUIRE(scheduler.schedule_graph(strided_imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & strided_dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(strided_dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_strided_f32"); + require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.hidden_size", "256"); + require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.input_stride", "288"); + require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.op", "2"); + REQUIRE(strided_dispatch.bindings.size() == 3); + REQUIRE(strided_dispatch.bindings[0].length == + static_cast(token_count - 1) * view->nb[1] + hidden_size * sizeof(float)); + + ggml_tensor * full_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * rejected_norm = ggml_rms_norm(ctx, view, 0.00001f); + ggml_tensor * rejected_scaled = ggml_mul(ctx, rejected_norm, full_weight); + REQUIRE(full_weight != nullptr); + REQUIRE(rejected_scaled != nullptr); + ggml_cgraph * rejected_graph = ggml_new_graph(ctx); + REQUIRE(rejected_graph != nullptr); + ggml_build_forward_expand(rejected_graph, rejected_scaled); + ggml::hrx::GraphImportResult rejected_imported = ggml::hrx::import_ggml_graph(*rejected_graph); + REQUIRE(rejected_imported.valid()); + REQUIRE(scheduler.schedule_graph(rejected_imported.graph, test_dispatch_target())); + for (const ggml::hrx::Dispatch & rejected_dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(rejected_dispatch.kernel.kernel_id) != + "loom_libs:ggml_rmsnorm_binary_strided_f32"); + } + + ggml_tensor * second_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * shared_norm = ggml_rms_norm(ctx, view, 0.00001f); + ggml_tensor * scaled_a = ggml_mul(ctx, shared_norm, weight); + ggml_tensor * scaled_b = ggml_add(ctx, shared_norm, second_weight); + ggml_tensor * joined = ggml_add(ctx, scaled_a, scaled_b); + REQUIRE(second_weight != nullptr); + REQUIRE(joined != nullptr); + ggml_cgraph * fanout_graph = ggml_new_graph(ctx); + REQUIRE(fanout_graph != nullptr); + ggml_build_forward_expand(fanout_graph, joined); + ggml::hrx::GraphImportResult fanout_imported = ggml::hrx::import_ggml_graph(*fanout_graph); + REQUIRE(fanout_imported.valid()); + REQUIRE(scheduler.schedule_graph(fanout_imported.graph, test_dispatch_target())); + for (const ggml::hrx::Dispatch & fanout_dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(fanout_dispatch.kernel.kernel_id) != + "loom_libs:ggml_rmsnorm_binary_strided_f32"); + } + } + + ggml_free(ctx); +} + +static void run_rope_set_rows_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); + + constexpr int64_t head_size = 8; + constexpr int64_t head_count = 2; + constexpr int64_t token_count = 3; + constexpr int64_t cache_rows = 8; + constexpr int64_t hidden_size = head_size * head_count; + + const auto check_rope_concat = [&](int64_t concat_head_count, int64_t concat_token_count, int64_t n_dims) { + constexpr int64_t prefix_size = 64; + constexpr int64_t rope_size = 32; + constexpr int64_t source_size = prefix_size + rope_size; + const size_t source_elements = static_cast(source_size) * concat_head_count * concat_token_count; + ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, source_elements); + ggml_tensor * prefix = + ggml_view_3d(ctx, storage, prefix_size, concat_head_count, concat_token_count, source_size * sizeof(float), + source_size * concat_head_count * sizeof(float), 0); + ggml_tensor * rope_input = + ggml_view_3d(ctx, storage, rope_size, concat_head_count, concat_token_count, source_size * sizeof(float), + source_size * concat_head_count * sizeof(float), prefix_size * sizeof(float)); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, concat_token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); + ggml_tensor * rope = ggml_rope_ext(ctx, rope_input, positions, freqs, n_dims, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * output = ggml_concat(ctx, prefix, rope, 0); + REQUIRE(storage != nullptr); + REQUIRE(prefix != nullptr); + REQUIRE(rope_input != nullptr); + REQUIRE(positions != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(rope != 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); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_concat_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == concat_token_count); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.prefix_size", std::to_string(prefix_size)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.rope_size", std::to_string(rope_size)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.n_dims", std::to_string(n_dims)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.head_count", std::to_string(concat_head_count)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.token_capacity", std::to_string(concat_token_count)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.input_stride1", std::to_string(source_size)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.input_stride2", + std::to_string(source_size * concat_head_count)); + require_compile_parameter(dispatch, "ggml.rope_concat_f32.mode", "2"); + REQUIRE(dispatch.bindings.size() == 5); + REQUIRE(dispatch.bindings[1].length == source_elements * sizeof(float)); + + 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)); + }; + + check_rope_concat(40, 64, 32); + check_rope_concat(3, 2, 16); + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * output = ggml_rope_ext(ctx, input, pos, freqs, head_size, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != 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()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ROPE); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + require_compile_parameter(dispatch, "ggml.rope_f32.head_size", "8"); + require_compile_parameter(dispatch, "ggml.rope_f32.n_dims", "8"); + require_compile_parameter(dispatch, "ggml.rope_f32.head_count", "2"); + require_compile_parameter(dispatch, "ggml.rope_f32.token_capacity", "3"); + require_compile_parameter(dispatch, "ggml.rope_f32.mode", "2"); + REQUIRE(dispatch.bindings.size() == 5); + + 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)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * output = ggml_rope_ext(ctx, input, pos, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, 10000.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != 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()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ROPE); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + require_compile_parameter(dispatch, "ggml.rope_f32.head_size", "8"); + require_compile_parameter(dispatch, "ggml.rope_f32.n_dims", "8"); + require_compile_parameter(dispatch, "ggml.rope_f32.head_count", "2"); + require_compile_parameter(dispatch, "ggml.rope_f32.token_capacity", "3"); + require_compile_parameter(dispatch, "ggml.rope_f32.mode", "0"); + REQUIRE(dispatch.bindings.size() == 5); + + 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)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * output = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); + REQUIRE(input != nullptr); + REQUIRE(pos != 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(scheduler.plan().transients.size() == 2); + REQUIRE(scheduler.plan().constant_initializations.size() == 2); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + require_compile_parameter(dispatch, "ggml.rope_f32.mode", "2"); + require_compile_parameter(dispatch, "ggml.rope_f32.n_dims", "8"); + } + + { + constexpr int64_t partial_head_size = 128; + constexpr int64_t partial_n_dims = 96; + constexpr int64_t partial_head_count = 24; + constexpr int64_t partial_tokens = 2; + ggml_tensor * input = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, partial_head_size, partial_head_count, partial_tokens); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, partial_tokens); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, partial_n_dims / 2); + ggml_tensor * output = ggml_rope_ext(ctx, input, pos, freqs, partial_n_dims, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, + 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != 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); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + require_compile_parameter(dispatch, "ggml.rope_f32.head_size", std::to_string(partial_head_size)); + require_compile_parameter(dispatch, "ggml.rope_f32.n_dims", std::to_string(partial_n_dims)); + require_compile_parameter(dispatch, "ggml.rope_f32.head_count", std::to_string(partial_head_count)); + require_compile_parameter(dispatch, "ggml.rope_f32.token_capacity", std::to_string(partial_tokens)); + require_compile_parameter(dispatch, "ggml.rope_f32.mode", "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)); + } + + { + constexpr int64_t partial_head_size = 128; + constexpr int64_t partial_n_dims = 96; + constexpr int64_t partial_head_count = 24; + constexpr int64_t partial_tokens = 2; + constexpr int64_t partial_token_stride = 5120; + constexpr size_t input_offset = 64; + constexpr size_t input_elements = input_offset / sizeof(float) + + static_cast(partial_tokens - 1) * partial_token_stride + + partial_head_size * partial_head_count; + + ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, input_elements); + ggml_tensor * input = + ggml_view_3d(ctx, storage, partial_head_size, partial_head_count, partial_tokens, + partial_head_size * sizeof(float), partial_token_stride * sizeof(float), input_offset); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, partial_tokens); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, partial_n_dims / 2); + ggml_tensor * output = ggml_rope_ext(ctx, input, pos, freqs, partial_n_dims, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, + 1.0f, 0.0f, 1.19024f, 32.0f, 1.0f); + REQUIRE(storage != nullptr); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != 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); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + require_compile_parameter(dispatch, "ggml.rope_f32.head_size", std::to_string(partial_head_size)); + require_compile_parameter(dispatch, "ggml.rope_f32.n_dims", std::to_string(partial_n_dims)); + require_compile_parameter(dispatch, "ggml.rope_f32.head_count", std::to_string(partial_head_count)); + require_compile_parameter(dispatch, "ggml.rope_f32.token_capacity", std::to_string(partial_tokens)); + require_compile_parameter(dispatch, "ggml.rope_f32.input_stride1", std::to_string(partial_head_size)); + require_compile_parameter(dispatch, "ggml.rope_f32.input_stride2", std::to_string(partial_token_stride)); + require_compile_parameter(dispatch, "ggml.rope_f32.mscale", "1.19024003"); + require_compile_parameter(dispatch, "ggml.rope_f32.mode", "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)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * raw_scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count); + ggml_tensor * scales = ggml_scale(ctx, raw_scales, 0.5f); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NORMAL); + ggml_tensor * output = ggml_mul(ctx, rope, scales); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(raw_scales != nullptr); + REQUIRE(scales != nullptr); + REQUIRE(rope != 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()); + const size_t rope_index = producer_index_for_tensor(imported.graph, rope); + const size_t scale_index = producer_index_for_tensor(imported.graph, scales); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch match; + + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, rope_index, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + + covered_nodes[scale_index] = true; + match = {}; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, rope_index, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_token_scale_f32"); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); + ggml_tensor * output = ggml_scale(ctx, rope, 0.125f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(rope != 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); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_scale_f32"); + require_compile_parameter(dispatch, "ggml.rope_scale_f32.scale", "0.125"); + REQUIRE(dispatch.bindings.size() == 5); + + 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)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NORMAL); + ggml_tensor * output = ggml_mul(ctx, rope, scales); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(scales != nullptr); + REQUIRE(rope != 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); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_token_scale_f32"); + REQUIRE(dispatch.bindings.size() == 6); + + 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)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); + ggml_tensor * output = ggml_scale(ctx, rope, 0.125f); + ggml_tensor * side = ggml_scale(ctx, rope, 0.25f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(output != nullptr); + REQUIRE(side != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + 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())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 3); + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_f32"); + } + + { + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(cache != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(indices != 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()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_SET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_set_rows"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_rows); + REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == hidden_size); + require_compile_parameter(dispatch, "ggml.set_rows.input_format", "32"); + require_compile_parameter(dispatch, "ggml.set_rows.output_format", "16"); + REQUIRE(dispatch.bindings.size() == 3); + + 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)); + } + + { + constexpr int64_t row_stride = hidden_size * 5; + constexpr size_t row_offset = 4 * sizeof(float); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); + ggml_tensor * storage = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, row_offset / sizeof(float) + (token_count - 1) * row_stride + hidden_size); + ggml_tensor * rows = + ggml_view_2d(ctx, storage, hidden_size, token_count, row_stride * sizeof(float), row_offset); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(cache != nullptr); + REQUIRE(storage != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(indices != 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()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_VIEW); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_SET_ROWS); + + const ggml::hrx::Value * rows_value = imported.graph.values().find(imported.graph.nodes()[1].inputs[0]); + REQUIRE(rows_value != nullptr); + const ggml::hrx::Value * storage_value = imported.graph.values().find(rows_value->storage_root); + REQUIRE(storage_value != nullptr); + REQUIRE(rows_value->storage_offset == row_offset); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_set_rows"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_rows); + REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == hidden_size); + require_compile_parameter(dispatch, "ggml.set_rows.input_format", "32"); + require_compile_parameter(dispatch, "ggml.set_rows.output_format", "16"); + require_compile_parameter(dispatch, "ggml.set_rows.input_stride", std::to_string(row_stride)); + REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(dispatch.bindings[0].value == storage_value->id); + REQUIRE(dispatch.bindings[0].offset == row_offset); + REQUIRE(dispatch.bindings[0].length == + static_cast(token_count - 1) * rows->nb[1] + static_cast(hidden_size) * sizeof(float)); + + 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)); + } + + { + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, token_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(cache != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(indices != 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); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_set_rows"); + require_compile_parameter(dispatch, "ggml.set_rows.input_format", "16"); + require_compile_parameter(dispatch, "ggml.set_rows.output_format", "16"); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * rope = ggml_rope_ext(ctx, input, pos, freqs, head_size, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * rows = ggml_reshape_2d(ctx, rope, hidden_size, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != 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()); + REQUIRE(imported.graph.nodes().size() == 3); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ROPE); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_RESHAPE); + REQUIRE(imported.graph.nodes()[2].op == GGML_OP_SET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_set_rows_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_rows); + require_compile_parameter(dispatch, "ggml.rope_set_rows_f32.output_format", "16"); + require_compile_parameter(dispatch, "ggml.rope_set_rows_f32.n_dims", "8"); + require_compile_parameter(dispatch, "ggml.rope_set_rows_f32.mode", "2"); + REQUIRE(dispatch.bindings.size() == 6); + + 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)); + } + + { + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * rope = ggml_rope_ext(ctx, input, pos, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, 10000.0f, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * rows = ggml_reshape_2d(ctx, rope, hidden_size, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != 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()); + REQUIRE(imported.graph.nodes().size() == 3); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ROPE); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_RESHAPE); + REQUIRE(imported.graph.nodes()[2].op == GGML_OP_SET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_set_rows_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_rows); + require_compile_parameter(dispatch, "ggml.rope_set_rows_f32.output_format", "16"); + require_compile_parameter(dispatch, "ggml.rope_set_rows_f32.n_dims", "8"); + require_compile_parameter(dispatch, "ggml.rope_set_rows_f32.mode", "0"); + REQUIRE(dispatch.bindings.size() == 6); + + 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)); + } + + ggml_free(ctx); +} + +static void run_graph_traversal_checks() { + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out0 = ggml_add(ctx, a, b); + ggml_tensor * out1 = ggml_add(ctx, c, d); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + ggml_build_forward_expand(graph, out1); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].output == imported.graph.values().find_tensor(out0)->id); + REQUIRE(imported.graph.nodes()[1].output == imported.graph.values().find_tensor(out1)->id); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 2); + REQUIRE(order[0] == 0); + REQUIRE(order[1] == 1); + + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * add_out = ggml_add(ctx, a, b); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 3); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 2); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(add_out != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(matmul != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, add_out); + ggml_build_forward_expand(graph, matmul); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 2); + REQUIRE(order[0] == 1); + REQUIRE(order[1] == 0); + + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 256 * 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, 256, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); + ggml_tensor * out = ggml_mul(ctx, rms, weight); + REQUIRE(input != nullptr); + REQUIRE(weight != nullptr); + REQUIRE(rms != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_RMS_NORM); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 2); + REQUIRE(order[0] == 0); + REQUIRE(order[1] == 1); + + ggml_free(ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * left = ggml_add(ctx, a, b); + ggml_tensor * right = ggml_mul(ctx, a, c); + ggml_tensor * join = ggml_add(ctx, left, right); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(left != nullptr); + REQUIRE(right != nullptr); + REQUIRE(join != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, join); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + + const size_t left_index = producer_index_for_tensor(imported.graph, left); + const size_t right_index = producer_index_for_tensor(imported.graph, right); + const size_t join_index = producer_index_for_tensor(imported.graph, join); + + const std::vector order = traversal_indices(imported.graph); + REQUIRE(order.size() == 3); + REQUIRE(find_position(order, join_index) > find_position(order, left_index)); + REQUIRE(find_position(order, join_index) > find_position(order, right_index)); + + ggml_free(ctx); + } +} + +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, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + const ggml_type weight_type = output->src[0]->type; + REQUIRE(!scheduler.plan().dispatches.empty()); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + const bool vector_kernel = string_contains(kernel_name, "ggml_mul_mat_vector"); + const bool q6_vector_kernel = kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"; + const bool 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"; + const bool q8_input = !vector_kernel && expected_token_count <= 5 && expected_output_size % 64 == 0 && + (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && q8_capable_kernel; + const size_t command_count = q8_input ? 2 : 1; + REQUIRE(scheduler.plan().dispatches.size() == command_count); + if (q8_input) { + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + } + + const bool accepts_vector_token1 = + expected_token_count == 1 && (expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || + expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" || + expected_kernel == "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma" || + expected_kernel == "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32"); + if (vector_kernel && accepts_vector_token1) { + REQUIRE(string_contains(kernel_name, "ggml_mul_mat_vector")); + } else { + REQUIRE(kernel_name == expected_kernel_name); + } + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + if (vector_kernel) { + 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)); + } + require_compile_parameter(dispatch, "ggml.matmul.vector.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.matmul.vector.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.matmul.vector.output_accumulation", "0"); + if (q6_vector_kernel) { + require_compile_parameter(dispatch, "ggml.matmul.vector.q6_storage", "0"); + } + } else if (string_contains(kernel_name, "ggml_mul_mat_q6_k_packed_token1_f16_wmma")) { + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.output_size", + std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.output_accumulation", "0"); + } else if (string_contains(kernel_name, "_decode")) { + REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(dispatch.kernel.integer_parameters.at("input_size") == expected_input_size); + REQUIRE(dispatch.kernel.integer_parameters.at("output_size") == expected_output_size); + require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.token_capacity", + std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.output_capacity", + std::to_string(expected_output_size)); + const ggml::hrx::Value * weight = imported.graph.values().find(dispatch.bindings[1].value); + REQUIRE(weight != nullptr); + 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() == (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"); + require_compile_parameter(dispatch, "ggml.mul_mat.output_unary_op", + std::to_string(ggml::hrx::unary_kind_config_value(ggml::hrx::UnaryKind::Identity))); + const ggml::hrx::Value * weight = imported.graph.values().find(dispatch.bindings[1].value); + REQUIRE(weight != nullptr); + require_compile_parameter(dispatch, "ggml.mul_mat.weight_format", + std::to_string(q8_input ? (weight_type == GGML_TYPE_Q4_K ? 44 : 46) : + matmul_weight_format_config(weight->type))); + } else if (string_contains(kernel_name, "dense_linear")) { + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.input_size", + std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_size", + std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_accumulation", "0"); + } else { + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expected_output_size)); + } + + 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() == command_count); + REQUIRE(command_program_verifies(commands)); + 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"); + 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[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, + ggml_tensor * output, + const char * expected_kernel_name, + ggml::hrx::UnaryKind expected_unary_op, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + const ggml_type weight_type = output->src[0]->src[0]->type; + const bool q8_input = expected_token_count <= 5 && expected_output_size % 64 == 0 && + (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K); + const size_t command_count = q8_input ? 2 : 1; + REQUIRE(scheduler.plan().dispatches.size() == command_count); + if (q8_input) { + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + } + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.bindings.size() == 3); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "ggml.mul_mat.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.mul_mat.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.mul_mat.output_accumulation", "0"); + require_compile_parameter(dispatch, "ggml.mul_mat.output_unary_op", + std::to_string(ggml::hrx::unary_kind_config_value(expected_unary_op))); + const ggml::hrx::Value * weight = imported.graph.values().find(dispatch.bindings[1].value); + REQUIRE(weight != nullptr); + require_compile_parameter(dispatch, "ggml.mul_mat.weight_format", + std::to_string(q8_input ? (weight_type == GGML_TYPE_Q4_K ? 44 : 46) : + matmul_weight_format_config(weight->type))); + REQUIRE(dispatch.bindings[2].value == imported.graph.nodes()[1].output); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == command_count); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 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"); +} + +static void schedule_fused_matmul_swiglu_command( + ggml_context * ctx, + ggml_tensor * output, + ggml_type expected_gate_type, + ggml_type expected_up_type, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size, + ggml::hrx::BinaryKind expected_binary_op = ggml::hrx::BinaryKind::SwiGLU, + ggml_op expected_output_op = GGML_OP_GLU) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); + REQUIRE(imported.graph.nodes()[2].op == expected_output_op); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + const bool low_token = expected_token_count <= 5 && expected_input_size % 256 == 0; + const bool q8_input = low_token && expected_gate_type == GGML_TYPE_Q4_K && expected_up_type == GGML_TYPE_Q4_K && + expected_output_size % 64 == 0; + const size_t command_count = q8_input ? 2 : 1; + REQUIRE(scheduler.plan().dispatches.size() == command_count); + if (q8_input) { + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + } + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == (q8_input ? "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot" : + low_token ? "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot" : + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32")); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.gate_weight_format", + std::to_string(q8_input ? 44 : matmul_weight_format_config(expected_gate_type))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.up_weight_format", + std::to_string(q8_input ? 44 : matmul_weight_format_config(expected_up_type))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", + std::to_string(ggml::hrx::binary_kind_config_value(expected_binary_op))); + REQUIRE(dispatch.bindings[3].value == imported.graph.nodes()[2].output); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == command_count); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 4); + REQUIRE(commands.commands.back().bindings[0].name == "input"); + REQUIRE(commands.commands.back().bindings[1].name == "gate_weight"); + REQUIRE(commands.commands.back().bindings[2].name == "up_weight"); + REQUIRE(commands.commands.back().bindings[3].name == "output"); +} + +static void schedule_fused_matmul_swiglu_postops_command(ggml_context * ctx, + ggml_tensor * output, + const char * expected_kernel_name, + size_t expected_node_count, + const std::vector & expected_ops, + const std::vector expected_binding_names) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == expected_node_count); + REQUIRE(expected_ops.size() == expected_node_count); + for (size_t i = 0; i < expected_ops.size(); ++i) { + REQUIRE(imported.graph.nodes()[i].op == expected_ops[i]); + } + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == 8); + REQUIRE(dispatch.bindings.size() == expected_binding_names.size()); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", "8"); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.input_size", "640"); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.output_size", "256"); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.gate_weight_format", + std::to_string(matmul_weight_format_config(GGML_TYPE_IQ4_NL))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.up_weight_format", + std::to_string(matmul_weight_format_config(GGML_TYPE_IQ4_NL))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", + std::to_string(ggml::hrx::binary_kind_config_value(ggml::hrx::BinaryKind::SwiGLU))); + REQUIRE(dispatch.bindings.back().value == imported.graph.nodes().back().output); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == expected_binding_names.size()); + for (size_t i = 0; i < expected_binding_names.size(); ++i) { + REQUIRE(commands.commands.front().bindings[i].name == expected_binding_names[i]); + } +} + +static void schedule_packed_glu_command(ggml_context * ctx, + ggml_tensor * output, + ggml::hrx::BinaryKind expected_binary_op, + int64_t expected_hidden_size, + int64_t expected_token_count, + bool expected_swapped) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GLU); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "loom_libs:ggml_binary_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == expected_hidden_size * expected_token_count); + require_compile_parameter(dispatch, "ggml.binary_f32.op", + std::to_string(ggml::hrx::binary_kind_config_value(expected_binary_op))); + require_compile_parameter(dispatch, "ggml.binary_f32.ne0", std::to_string(expected_hidden_size)); + require_compile_parameter(dispatch, "ggml.binary_f32.ne1", std::to_string(expected_token_count)); + REQUIRE(dispatch.bindings.size() == 3); + + const size_t half_offset = static_cast(expected_hidden_size) * sizeof(float); + REQUIRE(dispatch.bindings[0].offset == (expected_swapped ? half_offset : 0)); + REQUIRE(dispatch.bindings[1].offset == (expected_swapped ? 0 : half_offset)); + REQUIRE(dispatch.bindings[2].value == imported.graph.nodes()[0].output); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 3); + REQUIRE(commands.commands.front().bindings[0].name == "lhs"); + REQUIRE(commands.commands.front().bindings[1].name == "rhs"); + REQUIRE(commands.commands.front().bindings[2].name == "output"); +} + +static void schedule_packed_matmul_glu_command(ggml_context * ctx, + ggml_tensor * output, + ggml_type expected_weight_type, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size, + ggml::hrx::BinaryKind expected_binary_op, + bool expected_swapped) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_GLU); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.gate_weight_format", + std::to_string(matmul_weight_format_config(expected_weight_type))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.up_weight_format", + std::to_string(matmul_weight_format_config(expected_weight_type))); + require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", + std::to_string(ggml::hrx::binary_kind_config_value(expected_binary_op))); + + const size_t half_offset = + ggml_row_size(expected_weight_type, expected_input_size) * static_cast(expected_output_size); + REQUIRE(dispatch.bindings[1].offset == (expected_swapped ? half_offset : 0)); + REQUIRE(dispatch.bindings[2].offset == (expected_swapped ? 0 : half_offset)); + REQUIRE(dispatch.bindings[3].value == imported.graph.nodes()[1].output); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 4); + REQUIRE(commands.commands.front().bindings[0].name == "input"); + REQUIRE(commands.commands.front().bindings[1].name == "gate_weight"); + REQUIRE(commands.commands.front().bindings[2].name == "up_weight"); + REQUIRE(commands.commands.front().bindings[3].name == "output"); +} + +static void schedule_fused_matmul_postops_command(ggml_context * ctx, + ggml_tensor * output, + const char * expected_kernel_name, + ggml_type expected_weight_type, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size, + size_t expected_node_count, + const std::vector & expected_ops, + const std::vector expected_binding_names, + bool expected_rmsnorm) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == expected_node_count); + REQUIRE(expected_ops.size() == expected_node_count); + for (size_t i = 0; i < expected_ops.size(); ++i) { + REQUIRE(imported.graph.nodes()[i].op == expected_ops[i]); + } + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == (expected_rmsnorm ? 2 : 1)); + REQUIRE(scheduler.plan().completion_counter_requests.empty()); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.bindings.size() == expected_binding_names.size()); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "ggml.mul_mat_postops.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_postops.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "ggml.mul_mat_postops.weight_format", + std::to_string(matmul_weight_format_config(expected_weight_type))); + REQUIRE(dispatch.kernel.compile_parameters.count("ggml.mul_mat_postops.rms_epsilon") == 0); + if (expected_rmsnorm) { + const ggml::hrx::Dispatch & rmsnorm_dispatch = scheduler.plan().dispatches[1]; + REQUIRE(kernel_name_for_id(rmsnorm_dispatch.kernel.kernel_id) == + expected_plain_rmsnorm_binary_kernel(expected_output_size, expected_token_count)); + } + + 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() == (expected_rmsnorm ? 2 : 1)); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == expected_binding_names.size()); + for (size_t i = 0; i < expected_binding_names.size(); ++i) { + REQUIRE(commands.commands.front().bindings[i].name == expected_binding_names[i]); + } +} + +static void schedule_fused_llama_attention_matmul_command(ggml_context * ctx, + ggml_tensor * output, + const char * expected_kernel_name, + ggml_type expected_weight_type, + int64_t expected_token_count, + int64_t expected_input_size, + int64_t expected_output_size, + int64_t expected_head_size, + int64_t expected_head_count, + int64_t expected_cache_row_count, + int64_t expected_cache_format, + size_t expected_node_count, + const std::vector & expected_ops, + const std::vector expected_binding_names) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == expected_node_count); + REQUIRE(expected_ops.size() == expected_node_count); + for (size_t i = 0; i < expected_ops.size(); ++i) { + REQUIRE(imported.graph.nodes()[i].op == expected_ops[i]); + } + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.bindings.size() == expected_binding_names.size()); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); + require_compile_parameter(dispatch, "llm.attention_qkv.input_size", std::to_string(expected_input_size)); + require_compile_parameter(dispatch, "llm.attention_qkv.output_size", std::to_string(expected_output_size)); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", + std::to_string(matmul_weight_format_config(expected_weight_type))); + if (expected_head_size > 0) { + require_compile_parameter(dispatch, "llm.attention_qkv.head_size", std::to_string(expected_head_size)); + require_compile_parameter(dispatch, "llm.attention_qkv.head_count", std::to_string(expected_head_count)); + } + if (expected_cache_row_count > 0) { + require_compile_parameter(dispatch, "llm.attention_qkv.cache_row_count", + std::to_string(expected_cache_row_count)); + require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", + std::to_string(expected_cache_format)); + } + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == expected_binding_names.size()); + for (size_t i = 0; i < expected_binding_names.size(); ++i) { + REQUIRE(commands.commands.front().bindings[i].name == expected_binding_names[i]); + } +} + +static void require_matmul_swiglu_falls_back(ggml_context * ctx, ggml_tensor * output) { + 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); + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) != "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma"); + } +} + +static void require_matmul_swiglu_falls_back_to_compilable_plan(ggml_context * ctx, ggml_tensor * output) { + 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); + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) != "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma"); + } + + 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)); +} + +static void require_matmul_root_matches_identity_single(ggml_context * ctx, ggml_tensor * output) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(!imported.graph.nodes().empty()); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_MUL_MAT); + + const std::vector covered_nodes(imported.graph.nodes().size(), false); + const ggml::hrx::CommandPlan plan; + const ggml::hrx::DispatchMatchContext context = { + imported.graph, + &imported.graph.nodes()[0], + 0, + covered_nodes, + plan, + ggml::hrx::ValueId(static_cast(imported.graph.values().size())), + }; + ggml::hrx::DispatchMatch match; + REQUIRE(test_dispatch_registry().match(context, match)); + REQUIRE(match.covered_nodes.size() == 1); + REQUIRE(match.covered_nodes.front() == 0); + const ggml::hrx::Value * weight = imported.graph.values().find(imported.graph.nodes()[0].inputs[0]); + const ggml::hrx::Value * input = imported.graph.values().find(imported.graph.nodes()[0].inputs[1]); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + const bool q8_input = input->ne[1] <= 5 && weight->ne[1] % 64 == 0 && + (weight->type == GGML_TYPE_Q4_K || weight->type == GGML_TYPE_Q6_K); + REQUIRE(match.dispatches.size() == (q8_input ? 2 : 1)); + require_compile_parameter(match.dispatches.back(), "ggml.mul_mat.output_unary_op", + std::to_string(ggml::hrx::unary_kind_config_value(ggml::hrx::UnaryKind::Identity))); +} + +static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) { + 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()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); +} + +static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { + 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 * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); + ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + REQUIRE(input != nullptr); + REQUIRE(norm_weight != nullptr); + REQUIRE(vocab_weight != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); + ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); + ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, normalized); + REQUIRE(rms != nullptr); + REQUIRE(normalized != nullptr); + REQUIRE(logits != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, logits); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + + 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); + + const ggml::hrx::Dispatch & rms_dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(rms_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); + REQUIRE(rms_dispatch.kernel.integer_parameters.at("token_count") == token_count); + require_compile_parameter(rms_dispatch, "qwen3_moe.model.hidden_size", "2048"); + require_compile_parameter(rms_dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + REQUIRE(rms_dispatch.bindings.size() == 4); + + const ggml::hrx::CommandPlanTransient & q8_transient = scheduler.plan().transients.front(); + REQUIRE(q8_transient.size == qwen_q8_1_x4_size(token_count, 2048)); + REQUIRE(rms_dispatch.bindings[3].value == q8_transient.value); + REQUIRE(rms_dispatch.bindings[3].length == q8_transient.size); + + const ggml::hrx::Dispatch & vocab_dispatch = scheduler.plan().dispatches[1]; + REQUIRE(kernel_name_for_id(vocab_dispatch.kernel.kernel_id) == "qwen3_moe:ggml_linear_q6k_q8_1_x4"); + REQUIRE(vocab_dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(vocab_dispatch.kernel.integer_parameters.at("input_size") == 2048); + REQUIRE(vocab_dispatch.kernel.integer_parameters.at("output_size") == 151936); + require_compile_parameter(vocab_dispatch, "ggml.linear_q6k_q8_1_x4.token_capacity", std::to_string(token_count)); + require_compile_parameter(vocab_dispatch, "ggml.linear_q6k_q8_1_x4.output_capacity", "151936"); + REQUIRE(vocab_dispatch.bindings.size() == 3); + REQUIRE(vocab_dispatch.bindings[0].value == q8_transient.value); + REQUIRE(vocab_dispatch.bindings[0].length == q8_transient.size); + + 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() == 2); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[0].bindings.size() == 4); + REQUIRE(commands.commands[0].bindings[3].name == "q8_output"); + REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings.size() == 3); + REQUIRE(commands.commands[1].bindings[0].name == "q8_input"); + REQUIRE(commands.commands[1].bindings[0].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings[0].value == q8_transient.value); + + 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; + 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 * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + REQUIRE(input != nullptr); + REQUIRE(vocab_weight != nullptr); + ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, input); + REQUIRE(logits != nullptr); + + schedule_single_matmul_command(ctx, logits, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 2048, 151936); + ggml_free(ctx); +} + +static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_type) { + 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 * embedding_weight = ggml_new_tensor_2d(ctx, embedding_weight_type, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + REQUIRE(embedding_weight != nullptr); + REQUIRE(token_ids != nullptr); + REQUIRE(vocab_weight != nullptr); + + ggml_tensor * embedding = ggml_get_rows(ctx, embedding_weight, token_ids); + ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, embedding); + REQUIRE(embedding != nullptr); + REQUIRE(logits != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, logits); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); + + 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); + + const size_t q8_byte_count = qwen_q8_1_x4_size(1, 2048); + + const ggml::hrx::Dispatch & get_rows_dispatch = scheduler.plan().dispatches[0]; + REQUIRE(kernel_name_for_id(get_rows_dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_f32_next"); + REQUIRE(get_rows_dispatch.kernel.integer_parameters.at("token_count") == 1); + REQUIRE(get_rows_dispatch.kernel.integer_parameters.at("row_count") == 151936); + REQUIRE(get_rows_dispatch.kernel.integer_parameters.at("hidden_size") == 2048); + require_compile_parameter(get_rows_dispatch, "ggml.get_rows_f32.weight_format", + std::to_string(matmul_weight_format_config(embedding_weight_type))); + require_compile_parameter(get_rows_dispatch, "ggml.get_rows_f32.next_format", "81"); + REQUIRE(get_rows_dispatch.bindings.size() == 4); + + const ggml::hrx::CommandPlanTransient & q8_transient = scheduler.plan().transients.front(); + REQUIRE(q8_transient.size == q8_byte_count); + REQUIRE(get_rows_dispatch.bindings[3].value == q8_transient.value); + REQUIRE(get_rows_dispatch.bindings[3].length == q8_transient.size); + + const ggml::hrx::CommandPlanAlternateValue * alternate = + scheduler.plan().metadata.find_alternate_value(imported.graph.nodes()[0].output, GGML_TYPE_Q8_1, q8_byte_count); + REQUIRE(alternate != nullptr); + REQUIRE(alternate->alternate_value == q8_transient.value); + + const ggml::hrx::Dispatch & vocab_dispatch = scheduler.plan().dispatches[1]; + const std::string vocab_kernel_name = kernel_name_for_id(vocab_dispatch.kernel.kernel_id); + REQUIRE(vocab_dispatch.bindings.size() == 3); + if (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4") { + REQUIRE(vocab_dispatch.bindings[0].value == alternate->alternate_value); + REQUIRE(vocab_dispatch.bindings[0].length == alternate->byte_count); + } else { + REQUIRE(vocab_kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"); + REQUIRE(vocab_dispatch.bindings[0].value == imported.graph.nodes()[0].output); + require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.input_size", "2048"); + require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.output_size", "151936"); + require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.q6_storage", "0"); + } + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 2); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[0].bindings.size() == 4); + REQUIRE(commands.commands[0].bindings[3].name == "next_output"); + REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[1].bindings.size() == 3); + REQUIRE(commands.commands[1].bindings[0].name == + (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4" ? "q8_input" : "input")); + + ggml_free(ctx); +} + +static bool graph_is_supported(ggml_context * ctx, ggml_tensor * output) { + 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()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); +} + +struct ManualQwenRouterTop8Graph { + ggml::hrx::Graph graph; + ggml::hrx::ValueId route_ids; +}; + +static ggml::hrx::ValueId add_manual_graph_tensor(ggml::hrx::Graph & graph, + ggml_tensor * tensor, + ggml::hrx::ValueKind kind) { REQUIRE(tensor != nullptr); return graph.values().get_or_add_tensor_value(tensor, kind); } -static ManualQwenRouterTop8Graph build_manual_qwen_router_top8_graph(ggml_context * ctx, - int64_t expert_count, - int64_t route_count, - int64_t token_count) { - ManualQwenRouterTop8Graph manual; - ggml::hrx::Graph & graph = manual.graph; +static ManualQwenRouterTop8Graph build_manual_qwen_router_top8_graph(ggml_context * ctx, + int64_t expert_count, + int64_t route_count, + int64_t token_count) { + ManualQwenRouterTop8Graph manual; + ggml::hrx::Graph & graph = manual.graph; + + const ggml::hrx::ValueId logits = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, expert_count, token_count), ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId probs = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, expert_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId probs_reshaped = add_manual_graph_tensor( + graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, expert_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId argsort = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, expert_count, token_count), ggml::hrx::ValueKind::Transient); + manual.route_ids = add_manual_graph_tensor(graph, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, token_count), + ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId selected = add_manual_graph_tensor( + graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId weights_flat = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, route_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId sum = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId clamped = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId normalized = add_manual_graph_tensor( + graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, route_count, token_count), ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId route_weights = add_manual_graph_tensor( + graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, token_count), ggml::hrx::ValueKind::External); + + ggml::hrx::GraphNode & softmax = graph.add_node(GGML_OP_SOFT_MAX, probs, { logits }); + softmax.params = ggml::hrx::SoftMaxParams{ 1.0f, 0.0f }; + graph.add_node(GGML_OP_RESHAPE, probs_reshaped, { probs }); + ggml::hrx::GraphNode & argsort_node = graph.add_node(GGML_OP_ARGSORT, argsort, { probs }); + argsort_node.params = ggml::hrx::ArgsortParams{ GGML_SORT_ORDER_DESC }; + graph.add_node(GGML_OP_VIEW, manual.route_ids, { argsort }); + graph.add_node(GGML_OP_GET_ROWS, selected, { probs_reshaped, manual.route_ids }); + graph.add_node(GGML_OP_RESHAPE, weights_flat, { selected }); + graph.add_node(GGML_OP_SUM_ROWS, sum, { weights_flat }); + ggml::hrx::GraphNode & clamp = graph.add_node(GGML_OP_CLAMP, clamped, { sum }); + clamp.params = ggml::hrx::ClampParams{ 0.00006103515625f, std::numeric_limits::infinity() }; + graph.add_node(GGML_OP_DIV, normalized, { weights_flat, clamped }); + graph.add_node(GGML_OP_RESHAPE, route_weights, { normalized }); + REQUIRE(graph.build_index().success()); + return manual; +} + +static ggml::hrx::Graph build_manual_token_embedding_graph(ggml_tensor * weight, + ggml_tensor * token_ids, + ggml_tensor * output) { + ggml::hrx::Graph graph; + ggml::hrx::ValueId weight_value = graph.values().get_or_add_tensor_value(weight, ggml::hrx::ValueKind::External); + ggml::hrx::ValueId token_ids_value = + graph.values().get_or_add_tensor_value(token_ids, ggml::hrx::ValueKind::External); + ggml::hrx::ValueId output_value = graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); + graph.add_node(GGML_OP_GET_ROWS, output_value, { weight_value, token_ids_value }); + REQUIRE(graph.build_index().success()); + return graph; +} + +static bool manual_token_embedding_graph_is_supported(ggml_context * ctx, + ggml_type weight_type, + ggml_type token_ids_type, + ggml_type output_type, + int64_t hidden_size, + int64_t vocabulary_count, + int64_t token_count, + int64_t output_hidden_size = -1, + int64_t output_token_count = -1) { + if (output_hidden_size < 0) { + output_hidden_size = hidden_size; + } + if (output_token_count < 0) { + output_token_count = token_count; + } + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, hidden_size, vocabulary_count); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, token_ids_type, token_count); + ggml_tensor * output = ggml_new_tensor_2d(ctx, output_type, output_hidden_size, output_token_count); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + REQUIRE(output != nullptr); + + ggml::hrx::Graph graph = build_manual_token_embedding_graph(weight, token_ids, output); + return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +} + +static void schedule_token_embedding_command(ggml_context * ctx, + ggml_tensor * output, + int64_t expected_token_count, + int64_t expected_vocabulary_count, + int64_t expected_hidden_size, + const char * expected_weight_format) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "loom_libs:ggml_get_rows_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("row_count") == expected_vocabulary_count); + REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == expected_hidden_size); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", expected_weight_format); + REQUIRE(dispatch.bindings.size() == 3); + + const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 3); + REQUIRE(commands.commands.front().bindings[0].name == "token_ids"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "output"); +} + +static void run_qwen_token_embedding_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); + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 1, 151936, 2048, "4"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 13); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 13, 151936, 2048, "4"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 5, 151936, 2048, "30"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q1_0, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 1, 151936, 2048, "10"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_1, 640, 262144); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 1, 262144, 640, "51"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 3840, 262208); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 64); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 64, 262208, 3840, "80"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 3840, 262208); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 64); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 64, 262208, 3840, "6"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 3840, 262208); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 64); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 64, 262208, 3840, "30"); + } + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 262144, 4); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(weight != nullptr); + REQUIRE(token_ids != nullptr); + ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + REQUIRE(output != nullptr); + schedule_token_embedding_command(ctx, output, 1, 4, 262144, "32"); + } + + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F32, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q6_K, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q1_0, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q5_1, GGML_TYPE_I32, GGML_TYPE_F32, 640, 262144, 1)); + REQUIRE(manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ4_XS, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, + 1)); + REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ3_S, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, + 1)); + REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ4_NL, GGML_TYPE_I32, GGML_TYPE_F32, 2048, + 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q8_0, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q8_1, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F16, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_BF16, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_BF16, GGML_TYPE_I32, GGML_TYPE_F32, 3840, 262208, 64)); + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I64, GGML_TYPE_F32, 2048, 151936, 1)); + REQUIRE( + !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F16, 2048, 151936, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F32, GGML_TYPE_I32, GGML_TYPE_F32, 128, 151936, 1)); + REQUIRE(manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F32, GGML_TYPE_I32, GGML_TYPE_F32, 262144, 4, 1)); + REQUIRE( + manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F32, GGML_TYPE_I32, GGML_TYPE_F32, 33024, 151936, 1)); + REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, + 1, 2048, 2)); + + ggml_free(ctx); +} + +static void run_get_rows_rmsnorm_binary_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); + + ggml_tensor * embedding_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128256); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); + REQUIRE(embedding_weight != nullptr); + REQUIRE(token_ids != nullptr); + REQUIRE(norm_weight != nullptr); + ggml_tensor * embedding = ggml_get_rows(ctx, embedding_weight, token_ids); + ggml_tensor * rms = ggml_rms_norm(ctx, embedding, 0.00001f); + ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); + REQUIRE(embedding != nullptr); + REQUIRE(rms != nullptr); + REQUIRE(output != nullptr); + ggml_set_output(embedding); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 3); + + 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() == 1); + REQUIRE(plan.transients.size() == 2); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); + REQUIRE(dispatch.bindings.size() == 7); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "6"); + require_compile_parameter(dispatch, "ggml.get_rows_rmsnorm.rms_epsilon", "9.99999975e-06"); + + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); + REQUIRE(output_value != nullptr); + const size_t q8_bytes = qwen_q8_1_x4_size(512, 2048); + 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 == 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"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + + ggml_free(ctx); +} + +static void run_get_rows_scale_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); + + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); + REQUIRE(weight != nullptr); + REQUIRE(ids != nullptr); + ggml_tensor * rows = ggml_get_rows(ctx, weight, ids); + ggml_tensor * output = ggml_scale(ctx, rows, 0.177800179f); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + ggml_set_output(rows); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + require_scheduled_command_program( + graph, [&](const ggml::hrx::Graph & imported_graph, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(imported_graph.nodes().size() == 2); + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scale_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "32"); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.scale", expected_config_value(0.177800179f)); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.bias", expected_config_value(0.0f)); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings[0].name == "token_ids"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "raw_output"); + REQUIRE(commands.commands.front().bindings[3].name == "output"); + }); + + ggml_tensor * q5_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_1, 640, 262144); + ggml_tensor * q5_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 64); + REQUIRE(q5_weight != nullptr); + REQUIRE(q5_ids != nullptr); + ggml_tensor * q5_rows = ggml_get_rows(ctx, q5_weight, q5_ids); + ggml_tensor * q5_output = ggml_scale(ctx, q5_rows, 25.2982216f); + REQUIRE(q5_rows != nullptr); + REQUIRE(q5_output != nullptr); + ggml_set_output(q5_rows); + + ggml_cgraph * q5_graph = ggml_new_graph(ctx); + REQUIRE(q5_graph != nullptr); + ggml_build_forward_expand(q5_graph, q5_output); + require_scheduled_command_program( + q5_graph, [&](const ggml::hrx::Graph & imported_graph, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(imported_graph.nodes().size() == 2); + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scale_f32"); + REQUIRE(dispatch.bindings.size() == 4); + require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "51"); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.scale", expected_config_value(25.2982216f)); + require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.bias", expected_config_value(0.0f)); + REQUIRE(commands.commands.size() == 1); + }); + + ggml_tensor * q5_decode_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + REQUIRE(q5_decode_ids != nullptr); + ggml_tensor * q5_decode_rows = ggml_get_rows(ctx, q5_weight, q5_decode_ids); + ggml_tensor * q5_decode_output = ggml_scale(ctx, q5_decode_rows, 25.2982216f); + REQUIRE(q5_decode_rows != nullptr); + REQUIRE(q5_decode_output != nullptr); + ggml_cgraph * q5_decode_graph = ggml_new_graph(ctx); + REQUIRE(q5_decode_graph != nullptr); + ggml_build_forward_expand(q5_decode_graph, q5_decode_output); + const std::vector q5_decode_kernels = scheduled_kernel_names(q5_decode_graph); + REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == + q5_decode_kernels.end()); + REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_get_rows_f32") != + q5_decode_kernels.end()); + REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_scale_f32") != + q5_decode_kernels.end()); + + ggml_tensor * q4_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, 640, 64); + ggml_tensor * q4_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); + REQUIRE(q4_weight != nullptr); + REQUIRE(q4_ids != nullptr); + ggml_tensor * q4_rows = ggml_get_rows(ctx, q4_weight, q4_ids); + ggml_tensor * q4_output = ggml_scale(ctx, q4_rows, 0.5f); + REQUIRE(q4_rows != nullptr); + REQUIRE(q4_output != nullptr); + ggml_cgraph * q4_graph = ggml_new_graph(ctx); + REQUIRE(q4_graph != nullptr); + ggml_build_forward_expand(q4_graph, q4_output); + const std::vector q4_kernels = scheduled_kernel_names(q4_graph); + REQUIRE(std::find(q4_kernels.begin(), q4_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == q4_kernels.end()); + REQUIRE(std::find(q4_kernels.begin(), q4_kernels.end(), "loom_libs:ggml_get_rows_f32") != q4_kernels.end()); + + ggml_tensor * side_use = ggml_add(ctx, rows, rows); + REQUIRE(side_use != nullptr); + ggml_cgraph * fallback_graph = ggml_new_graph(ctx); + REQUIRE(fallback_graph != nullptr); + ggml_build_forward_expand(fallback_graph, output); + ggml_build_forward_expand(fallback_graph, side_use); + const std::vector fallback_kernels = scheduled_kernel_names(fallback_graph); + REQUIRE(std::find(fallback_kernels.begin(), fallback_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == + fallback_kernels.end()); + REQUIRE(std::find(fallback_kernels.begin(), fallback_kernels.end(), "loom_libs:ggml_get_rows_f32") != + fallback_kernels.end()); + + ggml_tensor * inplace_rows = ggml_get_rows(ctx, weight, ids); + ggml_tensor * inplace_output = ggml_scale_inplace(ctx, inplace_rows, 0.5f); + REQUIRE(inplace_rows != nullptr); + REQUIRE(inplace_output != nullptr); + ggml_cgraph * inplace_graph = ggml_new_graph(ctx); + REQUIRE(inplace_graph != nullptr); + ggml_build_forward_expand(inplace_graph, inplace_output); + const std::vector inplace_kernels = scheduled_kernel_names(inplace_graph); + REQUIRE(std::find(inplace_kernels.begin(), inplace_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == + inplace_kernels.end()); + REQUIRE(std::find(inplace_kernels.begin(), inplace_kernels.end(), "loom_libs:ggml_get_rows_f32") != + inplace_kernels.end()); + + ggml_free(ctx); +} + +static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, + int64_t hidden_size, + int64_t source_token_count, + int64_t output_token_count, + bool shared_row_ids, + int64_t second_source_hidden_size = -1) { + if (second_source_hidden_size < 0) { + second_source_hidden_size = hidden_size; + } + + ggml::hrx::Graph graph; + + ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, second_source_hidden_size, source_token_count); + ggml_tensor * row_ids0 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * row_ids1 = shared_row_ids ? row_ids0 : ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * selected0 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); + ggml_tensor * selected1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); + ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); + REQUIRE(attention != nullptr); + REQUIRE(residual != nullptr); + REQUIRE(row_ids0 != nullptr); + REQUIRE(row_ids1 != nullptr); + REQUIRE(selected0 != nullptr); + REQUIRE(selected1 != nullptr); + REQUIRE(output != nullptr); + + const ggml::hrx::ValueId attention_value = + graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId residual_value = + graph.values().get_or_add_tensor_value(residual, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids0_value = + graph.values().get_or_add_tensor_value(row_ids0, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids1_value = + graph.values().get_or_add_tensor_value(row_ids1, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId selected0_value = + graph.values().get_or_add_tensor_value(selected0, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId selected1_value = + graph.values().get_or_add_tensor_value(selected1, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId output_value = + graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); + + graph.add_node(GGML_OP_GET_ROWS, selected0_value, { attention_value, row_ids0_value }); + graph.add_node(GGML_OP_GET_ROWS, selected1_value, { residual_value, row_ids1_value }); + graph.add_node(GGML_OP_ADD, output_value, { selected0_value, selected1_value }); + REQUIRE(graph.build_index().success()); + return graph; +} + +struct GatherAddRmsNormGraph { + ggml_tensor * attention = nullptr; + ggml_tensor * residual = nullptr; + ggml_tensor * row_ids = nullptr; + ggml_tensor * weight = nullptr; + ggml_tensor * selected = nullptr; + ggml_tensor * raw_output = nullptr; + ggml_tensor * normalized_output = nullptr; + ggml_cgraph * graph = nullptr; +}; + +static GatherAddRmsNormGraph build_gather_add_rmsnorm_graph( + ggml_context * ctx, + int64_t hidden_size, + int64_t source_token_count, + int64_t output_token_count, + bool produced_row_ids = false, + bool produced_weight = false, + ggml::hrx::BinaryKind gather_op = ggml::hrx::BinaryKind::Add, + bool first_is_lhs = true) { + GatherAddRmsNormGraph result; + result.attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); + result.residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * raw_row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * raw_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + result.row_ids = produced_row_ids ? ggml_dup(ctx, raw_row_ids) : raw_row_ids; + result.weight = produced_weight ? ggml_scale(ctx, raw_weight, 0.5f) : raw_weight; + ggml_tensor * selected_attention = ggml_get_rows(ctx, result.attention, result.row_ids); + ggml_tensor * selected_residual = ggml_get_rows(ctx, result.residual, result.row_ids); + result.selected = selected_attention; + switch (gather_op) { + case ggml::hrx::BinaryKind::Add: + result.raw_output = ggml_add(ctx, selected_attention, selected_residual); + break; + case ggml::hrx::BinaryKind::Sub: + result.raw_output = first_is_lhs ? ggml_sub(ctx, selected_attention, selected_residual) : + ggml_sub(ctx, selected_residual, selected_attention); + break; + case ggml::hrx::BinaryKind::Mul: + result.raw_output = ggml_mul(ctx, selected_attention, selected_residual); + break; + case ggml::hrx::BinaryKind::Div: + result.raw_output = first_is_lhs ? ggml_div(ctx, selected_attention, selected_residual) : + ggml_div(ctx, selected_residual, selected_attention); + break; + default: + REQUIRE(false); + } + ggml_tensor * rms = ggml_rms_norm(ctx, result.raw_output, 0.000001f); + result.normalized_output = ggml_mul(ctx, rms, result.weight); + REQUIRE(result.attention != nullptr); + REQUIRE(result.residual != nullptr); + REQUIRE(raw_row_ids != nullptr); + REQUIRE(raw_weight != nullptr); + REQUIRE(result.row_ids != nullptr); + REQUIRE(result.weight != nullptr); + REQUIRE(selected_attention != nullptr); + REQUIRE(selected_residual != nullptr); + REQUIRE(result.raw_output != nullptr); + REQUIRE(rms != nullptr); + REQUIRE(result.normalized_output != nullptr); + + ggml_set_output(result.raw_output); + result.graph = ggml_new_graph(ctx); + REQUIRE(result.graph != nullptr); + ggml_build_forward_expand(result.graph, result.normalized_output); + return result; +} + +static ggml::hrx::DispatchRegistry gather_add_test_registry() { + ggml::hrx::DispatchRegistryBuilder builder; + ggml::hrx::register_gather_add_dispatch(builder); + return builder.build(); +} + +static bool match_gather_add_at(const ggml::hrx::DispatchRegistry & registry, + const ggml::hrx::Graph & graph, + size_t root_index, + const std::vector & covered_nodes, + ggml::hrx::DispatchMatch & match) { + ggml::hrx::CommandPlan plan; + const ggml::hrx::DispatchMatchContext context = { + graph, &graph.nodes()[root_index], + root_index, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + }; + return registry.match(context, match); +} + +static void require_gather_add_rmsnorm_route(ggml_context * ctx, + ggml::hrx::BinaryKind gather_op = ggml::hrx::BinaryKind::Add, + bool first_is_lhs = true) { + GatherAddRmsNormGraph tensors = + build_gather_add_rmsnorm_graph(ctx, 128, 4, 5, false, false, gather_op, first_is_lhs); + require_scheduled_command_program(tensors.graph, [&](const ggml::hrx::Graph &, const ggml::hrx::CommandPlan & plan, + const ggml::hrx::CommandProgram & commands) { + REQUIRE(plan.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_gather_add_rmsnorm_binary_f32"); + REQUIRE(dispatch.bindings.size() == 6); + require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.hidden_size", "128"); + require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", "9.99999997e-07"); + require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.gather_op", + std::to_string(ggml::hrx::binary_kind_config_value(gather_op))); + REQUIRE(commands.commands.size() == 1); + REQUIRE(commands.commands.front().bindings.size() == 6); + REQUIRE(commands.commands.front().bindings[0].name == "attention"); + REQUIRE(commands.commands.front().bindings[1].name == "residual"); + REQUIRE(commands.commands.front().bindings[2].name == "output_ids"); + REQUIRE(commands.commands.front().bindings[3].name == "raw_output"); + REQUIRE(commands.commands.front().bindings[4].name == "weight"); + REQUIRE(commands.commands.front().bindings[5].name == "normalized_output"); + }); +} + +static void require_gather_add_rmsnorm_availability(ggml_context * ctx) { + GatherAddRmsNormGraph tensors = build_gather_add_rmsnorm_graph(ctx, 128, 4, 5, true, true); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*tensors.graph); + REQUIRE(imported.valid()); + const ggml::hrx::DispatchRegistry registry = gather_add_test_registry(); + const size_t root_index = producer_index_for_tensor(imported.graph, tensors.selected); + const size_t row_index = producer_index_for_tensor(imported.graph, tensors.row_ids); + const size_t weight_index = producer_index_for_tensor(imported.graph, tensors.weight); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::DispatchMatch match; + + REQUIRE(!match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); + + covered_nodes[row_index] = true; + match = {}; + REQUIRE(match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); + + covered_nodes[weight_index] = true; + match = {}; + REQUIRE(match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == + "loom_libs:ggml_gather_add_rmsnorm_binary_f32"); +} + +static void require_gather_add_rmsnorm_row_ids_noalias(ggml_context * ctx) { + GatherAddRmsNormGraph tensors = build_gather_add_rmsnorm_graph(ctx, 128, 4, 5); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*tensors.graph); + REQUIRE(imported.valid()); + const ggml::hrx::Value * row_ids = imported.graph.values().find_tensor(tensors.row_ids); + const ggml::hrx::Value * attention = imported.graph.values().find_tensor(tensors.attention); + REQUIRE(row_ids != nullptr); + REQUIRE(attention != nullptr); + + ggml::hrx::Graph aliased_graph; + for (const ggml::hrx::ValueStorage & storage : imported.graph.values().storages()) { + REQUIRE(aliased_graph.values().add_snapshot_storage(storage).success()); + } + for (ggml::hrx::Value value : imported.graph.values().values()) { + if (value.id == row_ids->id) { + value.storage = attention->storage; + value.storage_root = attention->storage_root; + value.alias_source = attention->id; + value.storage_offset = attention->storage_offset; + value.storage_byte_count = attention->storage_byte_count; + } + REQUIRE(aliased_graph.values().add_snapshot_value(std::move(value)).success()); + } + for (const ggml::hrx::GraphNode & node : imported.graph.nodes()) { + ggml::hrx::GraphNode & copied = aliased_graph.add_node(node.op, node.output, node.inputs); + copied.params = node.params; + } + REQUIRE(aliased_graph.build_index().success()); + + const ggml::hrx::DispatchRegistry registry = gather_add_test_registry(); + size_t root_index = aliased_graph.nodes().size(); + for (size_t i = 0; i < aliased_graph.nodes().size(); ++i) { + if (aliased_graph.nodes()[i].op == GGML_OP_GET_ROWS) { + root_index = i; + break; + } + } + REQUIRE(root_index < aliased_graph.nodes().size()); + std::vector covered_nodes(aliased_graph.nodes().size(), false); + ggml::hrx::DispatchMatch match; + REQUIRE(match_gather_add_at(registry, aliased_graph, root_index, covered_nodes, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); +} + +static bool manual_gather_add_graph_is_supported(ggml_context * ctx, + int64_t hidden_size, + int64_t source_token_count, + int64_t output_token_count, + bool shared_row_ids, + int64_t second_source_hidden_size = -1) { + ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, hidden_size, source_token_count, output_token_count, + shared_row_ids, second_source_hidden_size); + return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +} + +static bool partial_gather_add_graph_is_supported(ggml_context * ctx) { + ggml::hrx::Graph graph; + + ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 13); + ggml_tensor * row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * add_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * add_output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(attention != nullptr); + REQUIRE(row_ids != nullptr); + REQUIRE(selected != nullptr); + REQUIRE(add_input != nullptr); + REQUIRE(add_output != nullptr); + + const ggml::hrx::ValueId attention_value = + graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids_value = + graph.values().get_or_add_tensor_value(row_ids, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId selected_value = + graph.values().get_or_add_tensor_value(selected, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId add_input_value = + graph.values().get_or_add_tensor_value(add_input, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId add_output_value = + graph.values().get_or_add_tensor_value(add_output, ggml::hrx::ValueKind::External); + + graph.add_node(GGML_OP_GET_ROWS, selected_value, { attention_value, row_ids_value }); + graph.add_node(GGML_OP_ADD, add_output_value, { selected_value, add_input_value }); + REQUIRE(graph.build_index().success()); + return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +} + +static void require_gather_add_rejects_unavailable_sources(ggml_context * ctx) { + ggml::hrx::Graph graph; + + ggml_tensor * raw_attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 13); + ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 13); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 13); + ggml_tensor * row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * selected0 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * selected1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(raw_attention != nullptr); + REQUIRE(attention != nullptr); + REQUIRE(residual != nullptr); + REQUIRE(row_ids != nullptr); + REQUIRE(selected0 != nullptr); + REQUIRE(selected1 != nullptr); + REQUIRE(output != nullptr); + + const ggml::hrx::ValueId raw_attention_value = + graph.values().get_or_add_tensor_value(raw_attention, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId attention_value = + graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId residual_value = + graph.values().get_or_add_tensor_value(residual, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId row_ids_value = + graph.values().get_or_add_tensor_value(row_ids, ggml::hrx::ValueKind::External); + const ggml::hrx::ValueId selected0_value = + graph.values().get_or_add_tensor_value(selected0, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId selected1_value = + graph.values().get_or_add_tensor_value(selected1, ggml::hrx::ValueKind::Transient); + const ggml::hrx::ValueId output_value = + graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); + + graph.add_node(GGML_OP_SCALE, attention_value, { raw_attention_value }); + graph.add_node(GGML_OP_GET_ROWS, selected0_value, { attention_value, row_ids_value }); + graph.add_node(GGML_OP_GET_ROWS, selected1_value, { residual_value, row_ids_value }); + graph.add_node(GGML_OP_ADD, output_value, { selected0_value, selected1_value }); + REQUIRE(graph.build_index().success()); + + ggml::hrx::DispatchRegistryBuilder builder; + ggml::hrx::register_gather_add_dispatch(builder); + const ggml::hrx::DispatchRegistry gather_add_registry = builder.build(); + + const size_t producer_index = 0; + const size_t get_rows_index = 1; + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch match; + std::vector covered_nodes(graph.nodes().size(), false); + ggml::hrx::DispatchMatchContext context = { + graph, &graph.nodes()[get_rows_index], + get_rows_index, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + }; + + REQUIRE(!gather_add_registry.match(context, match)); + + covered_nodes[producer_index] = true; + const ggml::hrx::DispatchMatchContext available_context = { + graph, &graph.nodes()[get_rows_index], + get_rows_index, covered_nodes, + plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + }; + match = {}; + REQUIRE(gather_add_registry.match(available_context, match)); + REQUIRE(match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); +} + +static void schedule_gather_add_command(ggml::hrx::Graph & graph, + int64_t expected_hidden_size, + int64_t expected_source_token_count, + int64_t expected_output_token_count) { + ggml::hrx::DispatchScheduler scheduler; + if (!scheduler.schedule_graph(graph, test_dispatch_target())) { + for (const std::string & error : scheduler.plan().status.errors()) { + std::fprintf(stderr, "%s\n", error.c_str()); + } + REQUIRE(false); + } + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 1); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "hrx:ggml_gather_add_f32"); + REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == expected_hidden_size); + REQUIRE(dispatch.kernel.integer_parameters.at("source_token_count") == expected_source_token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("output_token_count") == expected_output_token_count); + require_compile_parameter(dispatch, "ggml.gather_add_f32.binary_op", "0"); + REQUIRE(dispatch.bindings.size() == 4); + + const ggml::hrx::CommandProgram commands = + ggml::hrx::build_command_program(graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == 1); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 4); + REQUIRE(commands.commands.front().bindings[0].name == "attention"); + REQUIRE(commands.commands.front().bindings[1].name == "residual"); + REQUIRE(commands.commands.front().bindings[2].name == "output_ids"); + REQUIRE(commands.commands.front().bindings[3].name == "output"); +} + +static void run_gather_add_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); + + { + ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, 2048, 13, 1, true); + schedule_gather_add_command(graph, 2048, 13, 1); + } + { + ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, 2048, 128, 8, true); + schedule_gather_add_command(graph, 2048, 128, 8); + } + + REQUIRE(!manual_gather_add_graph_is_supported(ctx, 2048, 13, 1, false)); + REQUIRE(!manual_gather_add_graph_is_supported(ctx, 2048, 13, 1, true, 1024)); + REQUIRE(!manual_gather_add_graph_is_supported(ctx, 96, 13, 1, true)); + REQUIRE(!partial_gather_add_graph_is_supported(ctx)); + require_gather_add_rejects_unavailable_sources(ctx); + require_gather_add_rmsnorm_route(ctx); + require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Sub, false); + require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Mul); + require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Div); + require_gather_add_rmsnorm_availability(ctx); + require_gather_add_rmsnorm_row_ids_noalias(ctx); + + ggml_free(ctx); +} + +static void schedule_flash_attention_command(ggml_context * ctx, + ggml_tensor * output, + const char * expected_kernel_name, + const char * config_prefix, + 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), + bool transposed_value = false) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_FLASH_ATTN_EXT); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == (transposed_value ? 2 : 1)); + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("query_token_count") == query_token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); + REQUIRE(dispatch.bindings.size() == 6); + require_compile_parameter(dispatch, (std::string(config_prefix) + "query_head_count").c_str(), + std::to_string(query_head_count)); + require_compile_parameter(dispatch, (std::string(config_prefix) + "key_value_head_count").c_str(), + std::to_string(key_value_head_count)); + if (std::strcmp(config_prefix, "ggml.flash_attention.") == 0) { + 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.apply_gate", "0"); + require_compile_parameter(dispatch, "ggml.flash_attention.gate_stride_head", "1"); + require_compile_parameter(dispatch, "ggml.flash_attention.gate_stride_token", "1"); + } else { + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(query_token_count)); + } + + if (transposed_value) { + 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); + REQUIRE(copy.bindings[1].length == + static_cast(key_value_token_count * key_value_head_count * value_head_size) * + sizeof(ggml_fp16_t)); + require_compile_parameter(dispatch, "ggml.flash_attention.value_layout", "1"); + } else { + REQUIRE(dispatch.kernel.compile_parameters.count("ggml.flash_attention.value_layout") == 0); + } + + 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() == (transposed_value ? 2 : 1)); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 6); + REQUIRE(commands.commands.back().bindings[0].name == "query"); + REQUIRE(commands.commands.back().bindings[1].name == "key"); + REQUIRE(commands.commands.back().bindings[2].name == "value"); + REQUIRE(commands.commands.back().bindings[3].name == "mask"); + REQUIRE(commands.commands.back().bindings[4].name == "gate"); + REQUIRE(commands.commands.back().bindings[5].name == "output"); +} + +static void schedule_qwen_flash_attention_command(ggml_context * ctx, ggml_tensor * output) { + schedule_flash_attention_command(ctx, output, "loom_libs:ggml_flash_attention_f32_f16_wmma", + "ggml.flash_attention.", 16, 16, 4, 2); +} + +static void schedule_common_flash_attention_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)) { + schedule_flash_attention_command(ctx, output, "loom_libs:ggml_flash_attention_f32_f16_wmma", + "ggml.flash_attention.", query_token_count, key_value_token_count, + query_head_count, key_value_head_count, qk_head_size, value_head_size, scale); +} + +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); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_FLASH_ATTN_EXT); + + ggml::hrx::DispatchScheduler scheduler; + 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()); + + 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"); + REQUIRE(commands.valid()); + REQUIRE(commands.commands.size() == static_cast(query_token_count)); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.front().bindings.size() == 10); + REQUIRE(commands.commands.front().bindings[0].name == "query"); + 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[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); + const bool had_original_env = original_env != nullptr; + const std::string original_env_value = had_original_env ? original_env : ""; + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + for (const auto & shape : { + std::array{ 256, 512, 64, 1 }, + std::array{ 512, 512, 256, 1 }, + std::array{ 512, 2048, 320, 1 }, + std::array{ 255, 512, 256, 0 }, + std::array{ 512, 513, 256, 0 } + }) { + const float scale = 1.0f / std::sqrt(static_cast(shape[2])); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, shape[0], shape[1], 4, 2, GGML_TYPE_F32, + GGML_TYPE_F16, true, false, shape[2], scale); + schedule_flash_attention_command(ctx, output, "loom_libs:ggml_flash_attention_f32_f16_wmma", + "ggml.flash_attention.", shape[0], shape[1], 4, 2, shape[2], shape[2], scale, + shape[3] != 0); + } + + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 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); + schedule_qwen_flash_attention_command(ctx, output); + } + { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 16, 16, 4, 2); + schedule_common_flash_attention_command(ctx, output, 16, 16, 4, 2); + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + constexpr int64_t head_size = 256; + ggml_tensor * output = + build_qwen_flash_attention_graph(ctx, 16, 16, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, false, head_size, + 1.0f / std::sqrt(static_cast(head_size))); + schedule_common_flash_attention_command(ctx, output, 16, 16, 4, 2, head_size, head_size, + 1.0f / std::sqrt(static_cast(head_size))); + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + constexpr int64_t qk_head_size = 96; + constexpr int64_t value_head_size = 64; + ggml_tensor * output = build_qwen_flash_attention_graph( + ctx, 23, 256, 40, 40, 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_command(ctx, output, 23, 256, 40, 40, qk_head_size, value_head_size, + 1.0f / std::sqrt(static_cast(qk_head_size))); + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 8, 4, 2); + schedule_common_flash_attention_decode_fallback_command(ctx, output, 1, 8, 4, 2); + } + { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + constexpr int64_t qk_head_size = 96; + constexpr int64_t value_head_size = 64; + 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_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); + } + { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + constexpr int64_t tokens = 512; + constexpr int64_t hidden = 512; + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, tokens, tokens, 4, 2); + ggml_tensor * reshaped = ggml_reshape_2d(ctx, output, hidden, tokens); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, hidden, hidden); + ggml_tensor * projected = ggml_mul_mat(ctx, weight, reshaped); + ggml_cgraph * graph = ggml_new_graph(ctx); + 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()); + 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_flash_attention_f32_f16_wmma_publish_f16"; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings.size() == 7); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); + REQUIRE(output_value != nullptr); + const auto * alternate = + plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_F16, output_value->byte_count / 2); + REQUIRE(alternate != nullptr); + REQUIRE(producer->bindings.back().value == alternate->alternate_value); + const auto consumer = std::next(producer); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == alternate->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)); + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F32); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, false); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = + build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, true); + REQUIRE(!graph_is_supported(ctx, output)); + } + { + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 16, 16, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, + false, kQwenFlashHeadSize, 1.0f); + schedule_common_flash_attention_command(ctx, output, 16, 16, 4, 2, kQwenFlashHeadSize, kQwenFlashHeadSize, + 1.0f); + } + { + ggml_tensor * output = + build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, false, + kQwenFlashHeadSize, 1.0f / std::sqrt(128.0f), -1, 1.0f); + REQUIRE(!graph_is_supported(ctx, output)); + } + + ggml_free(ctx); + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); +} + +struct QwenAttentionPostprocessTensors { + ggml_tensor * query_raw = nullptr; + ggml_tensor * key_raw = nullptr; + ggml_tensor * value_raw = nullptr; + ggml_tensor * query_reshape = nullptr; + ggml_tensor * key_reshape = nullptr; + ggml_tensor * value_reshape = nullptr; + ggml_tensor * query_output = nullptr; + ggml_tensor * key_cache = nullptr; + ggml_tensor * value_cache = nullptr; + ggml_tensor * key_output = nullptr; + ggml_tensor * value_output = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * key_cache_indices = nullptr; + ggml_tensor * value_cache_indices = nullptr; + ggml_tensor * mask = nullptr; + ggml_tensor * flash_output = nullptr; +}; + +static QwenAttentionPostprocessTensors build_qwen_attention_postprocess_graph(ggml_context * ctx, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count, + float rms_epsilon = 0.000001f, + bool include_inverse_frequencies = true, + int rope_n_ctx_orig = 0, + float rope_freq_base = 10000.0f, + float rope_freq_scale = 1.0f, + float rope_ext_factor = 0.0f, + float rope_beta_fast = 0.0f, + float rope_beta_slow = 0.0f) { + QwenAttentionPostprocessTensors tensors; + const int64_t query_size = query_head_count * kQwenFlashHeadSize; + const int64_t key_value_size = key_value_head_count * kQwenFlashHeadSize; + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, token_count); + ggml_tensor * query_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, query_size); + ggml_tensor * key_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, key_value_size); + ggml_tensor * value_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, kQwenMoeHiddenSize, key_value_size); + REQUIRE(input != nullptr); + REQUIRE(query_weight != nullptr); + REQUIRE(key_weight != nullptr); + REQUIRE(value_weight != nullptr); + + tensors.query_raw = ggml_mul_mat(ctx, query_weight, input); + tensors.key_raw = ggml_mul_mat(ctx, key_weight, input); + tensors.value_raw = ggml_mul_mat(ctx, value_weight, input); + REQUIRE(tensors.query_raw != nullptr); + REQUIRE(tensors.key_raw != nullptr); + REQUIRE(tensors.value_raw != nullptr); + + tensors.query_reshape = ggml_reshape_3d(ctx, tensors.query_raw, kQwenFlashHeadSize, query_head_count, token_count); + tensors.key_reshape = ggml_reshape_3d(ctx, tensors.key_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + tensors.value_reshape = + ggml_reshape_3d(ctx, tensors.value_raw, kQwenFlashHeadSize, key_value_head_count, token_count); + REQUIRE(tensors.query_reshape != nullptr); + REQUIRE(tensors.key_reshape != nullptr); + REQUIRE(tensors.value_reshape != nullptr); + + ggml_tensor * query_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); + ggml_tensor * key_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); + tensors.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * inverse_frequencies = + include_inverse_frequencies ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize / 2) : nullptr; + REQUIRE(query_norm_weight != nullptr); + REQUIRE(key_norm_weight != nullptr); + REQUIRE(tensors.positions != nullptr); + REQUIRE(include_inverse_frequencies == (inverse_frequencies != nullptr)); + + ggml_tensor * query_norm = ggml_rms_norm(ctx, tensors.query_reshape, rms_epsilon); + ggml_tensor * query_mul = ggml_mul(ctx, query_norm, query_norm_weight); + REQUIRE(query_norm != nullptr); + REQUIRE(query_mul != nullptr); + tensors.query_output = ggml_rope_ext(ctx, query_mul, tensors.positions, inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, rope_n_ctx_orig, rope_freq_base, rope_freq_scale, + rope_ext_factor, 1.0f, rope_beta_fast, rope_beta_slow); + REQUIRE(tensors.query_output != nullptr); + + ggml_tensor * key_norm = ggml_rms_norm(ctx, tensors.key_reshape, rms_epsilon); + ggml_tensor * key_mul = ggml_mul(ctx, key_norm, key_norm_weight); + REQUIRE(key_norm != nullptr); + REQUIRE(key_mul != nullptr); + ggml_tensor * key_rope = ggml_rope_ext(ctx, key_mul, tensors.positions, inverse_frequencies, kQwenFlashHeadSize, + GGML_ROPE_TYPE_NEOX, rope_n_ctx_orig, rope_freq_base, rope_freq_scale, + rope_ext_factor, 1.0f, rope_beta_fast, rope_beta_slow); + REQUIRE(key_rope != nullptr); + + ggml_tensor * key_cache_rows = + ggml_reshape_2d(ctx, key_rope, kQwenFlashHeadSize * key_value_head_count, token_count); + ggml_tensor * value_cache_rows = + ggml_reshape_2d(ctx, tensors.value_reshape, kQwenFlashHeadSize * key_value_head_count, token_count); + REQUIRE(key_cache_rows != nullptr); + REQUIRE(value_cache_rows != nullptr); + + tensors.key_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); + tensors.value_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); + tensors.key_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + tensors.value_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(tensors.key_cache != nullptr); + REQUIRE(tensors.value_cache != nullptr); + REQUIRE(tensors.key_cache_indices != nullptr); + REQUIRE(tensors.value_cache_indices != nullptr); + + tensors.key_output = ggml_set_rows(ctx, tensors.key_cache, key_cache_rows, tensors.key_cache_indices); + tensors.value_output = ggml_set_rows(ctx, tensors.value_cache, value_cache_rows, tensors.value_cache_indices); + REQUIRE(tensors.key_output != nullptr); + REQUIRE(tensors.value_output != nullptr); + return tensors; +} + +static ggml::hrx::GraphImportResult import_qwen_attention_postprocess_graph( + ggml_context * ctx, + const QwenAttentionPostprocessTensors & tensors) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.query_output); + ggml_build_forward_expand(graph, tensors.key_output); + ggml_build_forward_expand(graph, tensors.value_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return imported; +} + +static ggml_tensor * append_qwen_flash_attention_consumer(ggml_context * ctx, + QwenAttentionPostprocessTensors & tensors, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count) { + ggml_tensor * query_layout = + ggml_reshape_3d(ctx, tensors.query_output, kQwenFlashHeadSize, query_head_count, token_count); + ggml_tensor * query_permute = ggml_permute(ctx, query_layout, 0, 2, 1, 3); + ggml_tensor * key_cache_layout = + ggml_reshape_3d(ctx, tensors.key_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); + ggml_tensor * key_permute = ggml_permute(ctx, key_cache_layout, 0, 2, 1, 3); + ggml_tensor * value_cache_layout = + ggml_reshape_3d(ctx, tensors.value_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); + ggml_tensor * value_permute = ggml_permute(ctx, value_cache_layout, 0, 2, 1, 3); + tensors.mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, cache_row_count, token_count); + REQUIRE(query_layout != nullptr); + REQUIRE(query_permute != nullptr); + REQUIRE(key_cache_layout != nullptr); + REQUIRE(key_permute != nullptr); + REQUIRE(value_cache_layout != nullptr); + REQUIRE(value_permute != nullptr); + REQUIRE(tensors.mask != nullptr); + + tensors.flash_output = ggml_flash_attn_ext(ctx, query_permute, key_permute, value_permute, tensors.mask, + 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)), 0.0f, 0.0f); + REQUIRE(tensors.flash_output != nullptr); + return tensors.flash_output; +} + +static void schedule_qwen_attention_postprocess_command(ggml_context * ctx, + const QwenAttentionPostprocessTensors & tensors, + int64_t token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t cache_row_count, + bool expect_synthetic_inverse_frequencies = false, + const std::string & expected_rope_mscale = "1") { + ggml::hrx::GraphImportResult imported = import_qwen_attention_postprocess_graph(ctx, tensors); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().initialization_dispatches.empty()); + const bool q8_input = token_count <= 5; + REQUIRE(scheduler.plan().dispatches.size() == (q8_input ? 5 : 4)); + if (q8_input) { + REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + } + + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_row_count); + REQUIRE(dispatch.bindings.size() == 12); + require_compile_parameter(dispatch, "qwen3_moe.model.rms_epsilon", "0.000001"); + require_compile_parameter(dispatch, "qwen3_moe.attention.head_size", std::to_string(kQwenFlashHeadSize)); + require_compile_parameter(dispatch, "qwen3_moe.attention.rope_mscale", expected_rope_mscale); + require_compile_parameter(dispatch, "qwen3_moe.attention.query_size", + std::to_string(query_head_count * kQwenFlashHeadSize)); + require_compile_parameter(dispatch, "qwen3_moe.attention.key_value_size", + std::to_string(key_value_head_count * kQwenFlashHeadSize)); + require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + + const ggml::hrx::Value * positions_value = imported.graph.values().find_tensor(tensors.positions); + const ggml::hrx::Value * key_indices_value = imported.graph.values().find_tensor(tensors.key_cache_indices); + const ggml::hrx::Value * value_indices_value = imported.graph.values().find_tensor(tensors.value_cache_indices); + const ggml::hrx::Value * query_raw_value = imported.graph.values().find_tensor(tensors.query_raw); + const ggml::hrx::Value * key_raw_value = imported.graph.values().find_tensor(tensors.key_raw); + const ggml::hrx::Value * value_raw_value = imported.graph.values().find_tensor(tensors.value_raw); + const ggml::hrx::Value * query_output_value = imported.graph.values().find_tensor(tensors.query_output); + const ggml::hrx::Value * key_cache_value = imported.graph.values().find_tensor(tensors.key_cache); + const ggml::hrx::Value * value_cache_value = imported.graph.values().find_tensor(tensors.value_cache); + REQUIRE(positions_value != nullptr); + REQUIRE(key_indices_value != nullptr); + REQUIRE(value_indices_value != nullptr); + REQUIRE(query_raw_value != nullptr); + REQUIRE(key_raw_value != nullptr); + REQUIRE(value_raw_value != nullptr); + REQUIRE(query_output_value != nullptr); + REQUIRE(key_cache_value != nullptr); + REQUIRE(value_cache_value != nullptr); + REQUIRE(dispatch.bindings[0].value == positions_value->id); + REQUIRE(dispatch.bindings[1].value == key_indices_value->id); + REQUIRE(dispatch.bindings[2].value == value_indices_value->id); + REQUIRE(dispatch.bindings[3].value == query_raw_value->id); + REQUIRE(dispatch.bindings[4].value == key_raw_value->id); + REQUIRE(dispatch.bindings[5].value == value_raw_value->id); + REQUIRE(dispatch.bindings[9].value == query_output_value->id); + REQUIRE(dispatch.bindings[10].value == key_cache_value->id); + REQUIRE(dispatch.bindings[11].value == value_cache_value->id); + if (expect_synthetic_inverse_frequencies) { + REQUIRE(scheduler.plan().transients.size() == (q8_input ? 2 : 1)); + REQUIRE(scheduler.plan().constant_initializations.size() == 1); + REQUIRE(dispatch.bindings[8].value == scheduler.plan().transients.back().value); + REQUIRE(scheduler.plan().constant_initializations[0].value == dispatch.bindings[8].value); + REQUIRE(scheduler.plan().constant_initializations[0].data.size() == + static_cast(kQwenFlashHeadSize / 2) * sizeof(float)); + } else { + REQUIRE(scheduler.plan().transients.size() == (q8_input ? 1 : 0)); + REQUIRE(scheduler.plan().constant_initializations.empty()); + } + + 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.initialization_commands.empty()); + REQUIRE(commands.commands.size() == (q8_input ? 5 : 4)); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands.back().bindings.size() == 12); + REQUIRE(commands.commands.back().bindings[0].name == "positions"); + REQUIRE(commands.commands.back().bindings[1].name == "key_cache_indices"); + REQUIRE(commands.commands.back().bindings[2].name == "value_cache_indices"); + REQUIRE(commands.commands.back().bindings[3].name == "query_input"); + REQUIRE(commands.commands.back().bindings[4].name == "key_input"); + REQUIRE(commands.commands.back().bindings[5].name == "value_input"); + REQUIRE(commands.commands.back().bindings[6].name == "query_norm_weight"); + REQUIRE(commands.commands.back().bindings[7].name == "key_norm_weight"); + REQUIRE(commands.commands.back().bindings[8].name == "inverse_frequencies"); + REQUIRE(commands.commands.back().bindings[9].name == "query_output"); + REQUIRE(commands.commands.back().bindings[10].name == "key_cache"); + REQUIRE(commands.commands.back().bindings[11].name == "value_cache"); + REQUIRE(commands.constant_initializations.size() == (expect_synthetic_inverse_frequencies ? 1 : 0)); +} + +static void run_qwen_attention_postprocess_dispatch_checks() { + 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 token_count = 4; + constexpr int64_t query_head_count = 4; + constexpr int64_t key_value_head_count = 2; + constexpr int64_t cache_row_count = 16; + const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count); + schedule_qwen_attention_postprocess_command(ctx, tensors, token_count, query_head_count, key_value_head_count, + cache_row_count); + } + + { + const QwenAttentionPostprocessTensors tensors = + build_qwen_attention_postprocess_graph(ctx, 4, 4, 2, 16, 0.00001f); + ggml::hrx::GraphImportResult imported = import_qwen_attention_postprocess_graph(ctx, tensors); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) != + "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16"); + } + } + + { + const QwenAttentionPostprocessTensors tensors = + build_qwen_attention_postprocess_graph(ctx, 4, 4, 2, 16, 0.000001f, false); + schedule_qwen_attention_postprocess_command(ctx, tensors, 4, 4, 2, 16, true); + } + + { + const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, 4, 4, 2, 16, 0.000001f, false, 8192, 1000000.0f, 0.25f, 1.0f, 32.0f, 1.0f); + schedule_qwen_attention_postprocess_command(ctx, tensors, 4, 4, 2, 16, true, "1.13862944"); + } + + { + constexpr int64_t token_count = 13; + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + constexpr int64_t cache_row_count = 512; + const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + schedule_qwen_attention_postprocess_command(ctx, tensors, token_count, query_head_count, key_value_head_count, + cache_row_count, true); + } + + { + constexpr int64_t token_count = 13; + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + constexpr int64_t cache_row_count = 512; + + const QwenAttentionPostprocessTensors first = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + const QwenAttentionPostprocessTensors second = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, first.query_output); + ggml_build_forward_expand(graph, first.key_output); + ggml_build_forward_expand(graph, first.value_output); + ggml_build_forward_expand(graph, second.query_output); + ggml_build_forward_expand(graph, second.key_output); + ggml_build_forward_expand(graph, second.value_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() == 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() == 8); + REQUIRE(command_program_verifies(commands)); + } + + { + constexpr int64_t token_count = 4; + constexpr int64_t query_head_count = 4; + constexpr int64_t key_value_head_count = 2; + constexpr int64_t cache_row_count = 16; + QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( + ctx, token_count, query_head_count, key_value_head_count, cache_row_count); + ggml_tensor * flash_output = append_qwen_flash_attention_consumer(ctx, tensors, token_count, query_head_count, + key_value_head_count, cache_row_count); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.key_output); + ggml_build_forward_expand(graph, tensors.value_output); + ggml_build_forward_expand(graph, flash_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, + producer_index_for_tensor(imported.graph, tensors.query_reshape), match)); + append_match_to_plan(plan, match, covered_nodes); + REQUIRE(plan.valid()); + REQUIRE(plan.initialization_dispatches.size() == 2); + + const ggml::hrx::Dispatch & context_capture = plan.initialization_dispatches[0]; + const ggml::hrx::Dispatch & metadata = plan.initialization_dispatches[1]; + REQUIRE(kernel_name_for_id(context_capture.kernel.kernel_id) == "qwen:qwen_attention_context_base_capture"); + REQUIRE(kernel_name_for_id(metadata.kernel.kernel_id) == "qwen3_moe:qwen_attention_metadata"); + REQUIRE(context_capture.bindings.size() == 2); + REQUIRE(metadata.bindings.size() == 5); + REQUIRE(context_capture.bindings[1].value == metadata.bindings[0].value); + REQUIRE(metadata.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(metadata.kernel.integer_parameters.at("context_capacity") == cache_row_count); + + const ggml::hrx::Value * positions_value = imported.graph.values().find_tensor(tensors.positions); + const ggml::hrx::Value * key_indices_value = imported.graph.values().find_tensor(tensors.key_cache_indices); + const ggml::hrx::Value * value_indices_value = imported.graph.values().find_tensor(tensors.value_cache_indices); + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(tensors.mask); + REQUIRE(positions_value != nullptr); + REQUIRE(key_indices_value != nullptr); + REQUIRE(value_indices_value != nullptr); + REQUIRE(mask_value != nullptr); + REQUIRE(context_capture.bindings[0].value == positions_value->id); + REQUIRE(metadata.bindings[1].value == positions_value->id); + REQUIRE(metadata.bindings[2].value == key_indices_value->id); + REQUIRE(metadata.bindings[3].value == value_indices_value->id); + REQUIRE(metadata.bindings[4].value == mask_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(commands.initialization_commands.size() == 2); + REQUIRE(commands.initialization_commands[0].bindings[0].name == "positions"); + REQUIRE(commands.initialization_commands[0].bindings[1].name == "control"); + REQUIRE(commands.initialization_commands[1].bindings[0].name == "control"); + REQUIRE(commands.initialization_commands[1].bindings[4].name == "attention_mask"); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, context_capture.bindings[1].value) != + nullptr); + + ggml::hrx::CommandProgram initialization_commands = copy_command_program_shape(commands); + initialization_commands.commands.clear(); + REQUIRE(command_program_verifies(initialization_commands)); + } + + ggml_free(ctx); +} + +static void run_packed_f16_producer_consumer_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + const struct { + ggml_type type; + int64_t tokens; + int64_t hidden; + int64_t outputs; + bool side_use; + bool packed; + bool alternate; + bool packed_copy = false; + } cases[] = { + { 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) { + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, test.tokens); + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 256, test.hidden); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 256, test.hidden); + 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 * weight = ggml_new_tensor_2d(ctx, test.type, test.hidden, test.outputs); + ggml_tensor * output = ggml_mul_mat(ctx, weight, glu); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + if (test.side_use) { + ggml_build_forward_expand(graph, ggml_scale(ctx, glu, 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 auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + const auto * value = imported.graph.values().find_tensor(glu); + REQUIRE(value != nullptr); + const auto * generated = + 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.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"; + }); + REQUIRE(producer != plan.dispatches.end()); + require_compile_parameter(*producer, "ggml.mul_mat_swiglu.f16_output_layout", test.packed ? "1" : "0"); + if (test.packed) { + REQUIRE(generated->byte_count == value->byte_count / 2); + REQUIRE(producer->bindings.back().value == generated->generated_value); + const auto consumer = std::next(producer); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == generated->generated_value); + if (test.type == GGML_TYPE_Q4_K) { + require_compile_parameter(*consumer, "ggml.mul_mat.f16_input_layout", "1"); + } + } else if (test.packed_copy) { + const auto copy = std::next(producer); + REQUIRE(copy != plan.dispatches.end()); + REQUIRE(kernel_name_for_id(copy->kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"); + REQUIRE(copy->bindings.front().value == alternate->alternate_value); + REQUIRE(copy->bindings.back().value == generated->generated_value); + REQUIRE(std::next(copy) != plan.dispatches.end()); + REQUIRE(std::next(copy)->bindings.front().value == generated->generated_value); + } + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + + ggml_free(ctx); +} + +static void run_generic_swiglu_k16_publish_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 output_size = 2048; + 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 * projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, hidden_size, output_size); + ggml_tensor * output = ggml_mul_mat(ctx, projection_weight, glu); + 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 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_k16"; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings.size() == 5); + const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(glu); + REQUIRE(glu_value != nullptr); + const auto * packed = + plan.metadata.find_generated_resource(glu_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + REQUIRE(packed != nullptr); + REQUIRE(packed->byte_count == glu_value->byte_count / 2); + REQUIRE(producer->bindings.back().value == packed->generated_value); + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*producer && !dispatch.bindings.empty() && + dispatch.bindings.front().value == packed->generated_value; + }); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(std::none_of(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"; + })); + + 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_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; + const char * producer_kernel; + const char * consumer_kernel; + ggml_type alternate_type; + }; + + const Case cases[] = { + { GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", + "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", GGML_TYPE_F16 }, + { GGML_TYPE_Q5_K, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", GGML_TYPE_Q8_1 }, + }; + + for (const Case & test : cases) { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); + ggml_tensor * producer_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 2048); + ggml_tensor * projection = ggml_mul_mat(ctx, producer_weight, input); + ggml_tensor * consumer_weight = ggml_new_tensor_2d(ctx, test.consumer_type, 2048, 1024); + ggml_tensor * output = ggml_mul_mat(ctx, consumer_weight, projection); + + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const auto * projection_value = imported.graph.values().find_tensor(projection); + REQUIRE(projection_value != nullptr); + const size_t alternate_bytes = test.alternate_type == GGML_TYPE_F16 ? + projection_value->byte_count / 2 : + static_cast(256) * ggml_row_size(GGML_TYPE_Q8_1, 2048); + const auto * alternate = + plan.metadata.find_alternate_value(projection_value->id, test.alternate_type, alternate_bytes); + REQUIRE(alternate != nullptr); + + const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == test.producer_kernel; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings[2].value == projection_value->id); + if (test.alternate_type == GGML_TYPE_F16) { + REQUIRE(producer->bindings.size() == 4); + REQUIRE(producer->bindings[3].value == alternate->alternate_value); + } else { + const auto quantize = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32" && + dispatch.bindings.front().value == projection_value->id; + }); + REQUIRE(quantize != plan.dispatches.end()); + REQUIRE(quantize->bindings.back().value == alternate->alternate_value); + } + + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == test.consumer_kernel; + }); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == alternate->alternate_value); + + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + ggml_free(ctx); + } + + { + 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() { + ggml_init_params params = {}; + params.mem_size = 16 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + auto build_glu = [&]() { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 8); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(glu != nullptr); + return glu; + }; + + { + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(bias != nullptr); + ggml_tensor * output = ggml_add(ctx, build_glu(), bias); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 4, + { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, + { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); + } - const ggml::hrx::ValueId logits = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, expert_count, token_count), ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId probs = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, expert_count, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId probs_reshaped = add_manual_graph_tensor( - graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, expert_count, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId argsort = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, expert_count, token_count), ggml::hrx::ValueKind::Transient); - manual.route_ids = add_manual_graph_tensor(graph, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, token_count), - ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId selected = add_manual_graph_tensor( - graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId weights_flat = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, route_count, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId sum = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId clamped = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId normalized = add_manual_graph_tensor( - graph, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, route_count, token_count), ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId route_weights = add_manual_graph_tensor( - graph, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, token_count), ggml::hrx::ValueKind::External); + { + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 8); + REQUIRE(residual != nullptr); + ggml_tensor * output = ggml_add(ctx, build_glu(), residual); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 4, + { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, + { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); + } - ggml::hrx::GraphNode & softmax = graph.add_node(GGML_OP_SOFT_MAX, probs, { logits }); - softmax.params = ggml::hrx::SoftMaxParams{ 1.0f, 0.0f }; - graph.add_node(GGML_OP_RESHAPE, probs_reshaped, { probs }); - ggml::hrx::GraphNode & argsort_node = graph.add_node(GGML_OP_ARGSORT, argsort, { probs }); - argsort_node.params = ggml::hrx::ArgsortParams{ GGML_SORT_ORDER_DESC }; - graph.add_node(GGML_OP_VIEW, manual.route_ids, { argsort }); - graph.add_node(GGML_OP_GET_ROWS, selected, { probs_reshaped, manual.route_ids }); - graph.add_node(GGML_OP_RESHAPE, weights_flat, { selected }); - graph.add_node(GGML_OP_SUM_ROWS, sum, { weights_flat }); - ggml::hrx::GraphNode & clamp = graph.add_node(GGML_OP_CLAMP, clamped, { sum }); - clamp.params = ggml::hrx::ClampParams{ 0.00006103515625f, std::numeric_limits::infinity() }; - graph.add_node(GGML_OP_DIV, normalized, { weights_flat, clamped }); - graph.add_node(GGML_OP_RESHAPE, route_weights, { normalized }); - REQUIRE(graph.build_index().success()); - return manual; -} + { + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 8); + REQUIRE(bias != nullptr); + REQUIRE(residual != nullptr); + ggml_tensor * biased = ggml_add(ctx, build_glu(), bias); + REQUIRE(biased != nullptr); + ggml_tensor * output = ggml_add(ctx, biased, residual); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 5, + { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD, GGML_OP_ADD }, + { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); + } -static ggml::hrx::Graph build_manual_token_embedding_graph(ggml_tensor * weight, - ggml_tensor * token_ids, - ggml_tensor * output) { - ggml::hrx::Graph graph; - ggml::hrx::ValueId weight_value = graph.values().get_or_add_tensor_value(weight, ggml::hrx::ValueKind::External); - ggml::hrx::ValueId token_ids_value = - graph.values().get_or_add_tensor_value(token_ids, ggml::hrx::ValueKind::External); - ggml::hrx::ValueId output_value = graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); - graph.add_node(GGML_OP_GET_ROWS, output_value, { weight_value, token_ids_value }); - REQUIRE(graph.build_index().success()); - return graph; + ggml_free(ctx); } -static bool manual_token_embedding_graph_is_supported(ggml_context * ctx, - ggml_type weight_type, - ggml_type token_ids_type, - ggml_type output_type, - int64_t hidden_size, - int64_t vocabulary_count, - int64_t token_count, - int64_t output_hidden_size = -1, - int64_t output_token_count = -1) { - if (output_hidden_size < 0) { - output_hidden_size = hidden_size; +static void run_swiglu_q8_output_checks() { + const struct { + int64_t tokens; + int64_t inputs; + int64_t outputs; + ggml_type down_type; + int projections; + bool terminal; + bool side_use; + bool observed; + bool geglu; + bool weight_alias; + bool packed; + } cases[] = { + { 2, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 3, 5120, 17408, GGML_TYPE_Q6_K, 1, false, false, false, false, false, true }, + { 4, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 5120, 17408, GGML_TYPE_Q6_K, 2, false, true, true, false, false, true }, + { 5, 4096, 16384, GGML_TYPE_Q4_K, 1, true, false, false, false, false, true }, + { 5, 5120, 32768, GGML_TYPE_Q6_K, 1, false, false, false, false, false, true }, + { 1, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 6, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, + { 5, 2048, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 6144, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 5120, 8192, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 5120, 33024, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, + { 5, 5120, 17408, GGML_TYPE_F32, 1, false, true, false, false, false, false }, + { 5, 5120, 17408, GGML_TYPE_Q6_K, 1, true, false, false, false, false, true }, + { 5, 5120, 17408, GGML_TYPE_Q4_K, 0, false, true, true, false, false, false }, + { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, true, false, false }, + { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, true, false }, + }; + + 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 * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.inputs, test.outputs); + ggml_tensor * up_weight = test.weight_alias ? + ggml_view_tensor(ctx, gate_weight) : + ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.inputs, test.outputs); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + ggml_tensor * glu = ggml_glu_split(ctx, gate, up, test.geglu ? GGML_GLU_OP_GEGLU : GGML_GLU_OP_SWIGLU); + if (test.observed) { + ggml_set_output(glu); + } + ggml_cgraph * graph = ggml_new_graph(ctx); + for (int i = 0; i < test.projections; ++i) { + ggml_tensor * down_weight = ggml_new_tensor_2d(ctx, test.down_type, test.outputs, 5120); + ggml_tensor * down = ggml_mul_mat(ctx, down_weight, glu); + ggml_build_forward_expand(graph, test.terminal ? down : ggml_add(ctx, down, ggml_dup_tensor(ctx, down))); + } + ggml_tensor * side = test.side_use ? ggml_scale(ctx, glu, 0.5f) : nullptr; + if (side != nullptr) { + ggml_build_forward_expand(graph, side); + } + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + const bool scheduled = scheduler.schedule_graph(imported.graph, test_dispatch_target()); + const auto & plan = scheduler.plan(); + 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_q8_1_x4_lowtoken_dot_q8_output"; + }); + REQUIRE((producer != plan.dispatches.end()) == test.packed); + if (test.outputs > 32768) { + REQUIRE(!scheduled); + REQUIRE(status_contains(plan.status, "unsupported HRX node")); + ggml_free(ctx); + continue; + } + REQUIRE(scheduled); + REQUIRE(plan.valid()); + const auto * output = imported.graph.values().find_tensor(glu); + REQUIRE(output != nullptr); + if (test.packed) { + const size_t bytes = static_cast(test.tokens) * ggml_row_size(GGML_TYPE_Q8_1, test.outputs); + const auto * alternate = plan.metadata.find_alternate_value(output->id, GGML_TYPE_Q8_1, bytes); + REQUIRE(alternate != nullptr); + REQUIRE(producer->bindings.size() == 5); + REQUIRE(producer->bindings[3].value == output->id); + REQUIRE(producer->bindings[4].value == alternate->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" && + dispatch.bindings.front().value == output->id; + }) == 0); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*producer && !dispatch.bindings.empty() && + dispatch.bindings.front().value == alternate->alternate_value; + }) == test.projections); + } + if (side != nullptr) { + const auto * side_value = imported.graph.values().find_tensor(side); + REQUIRE(side_value != nullptr); + const auto consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return !dispatch.bindings.empty() && dispatch.bindings.back().value == side_value->id; + }); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == output->id); + } + 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); } - if (output_token_count < 0) { - output_token_count = token_count; +} + +static void run_k16_major_preparation_reuse_checks() { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + const struct { + ggml_type first_type; + ggml_type second_type; + bool shared; + bool reshape; + } cases[] = { + { GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, true, false }, + { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, true, false }, + { GGML_TYPE_Q6_K, GGML_TYPE_Q4_K, true, false }, + { GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, true, false }, + { GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, false, false }, + { GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, false, false }, + { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, false, true }, + { GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, false, true }, + }; + + for (const auto & test : cases) { + ggml_tensor * first = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); + ggml_tensor * second = test.reshape ? ggml_reshape_2d(ctx, first, 2048, 1024) : + test.shared ? first : + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); + ggml_tensor * inputs[] = { first, second }; + ggml_tensor * outputs[2]; + ggml_cgraph * graph = ggml_new_graph(ctx); + for (int i = 0; i < 2; ++i) { + const ggml_type type = i == 0 ? test.first_type : test.second_type; + ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, inputs[i]->ne[0], 4096); + outputs[i] = ggml_mul_mat(ctx, weight, inputs[i]); + ggml_build_forward_expand(graph, outputs[i]); + } + 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 copy_count = + 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"; + }); + REQUIRE(copy_count == (test.shared ? 1 : 2)); + ggml::hrx::ValueId packed[2]; + for (int i = 0; i < 2; ++i) { + const auto * input = imported.graph.values().find_tensor(inputs[i]); + const auto * output = imported.graph.values().find_tensor(outputs[i]); + REQUIRE(input != nullptr); + REQUIRE(output != nullptr); + const auto * generated = + plan.metadata.find_generated_resource(input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + REQUIRE(generated != nullptr); + REQUIRE(generated->byte_count == input->byte_count / 2); + packed[i] = generated->generated_value; + const auto consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), + [&](const auto & dispatch) { return dispatch.bindings.back().value == output->id; }); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == packed[i]); + } + REQUIRE((packed[0] == packed[1]) == test.shared); + 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_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, hidden_size, vocabulary_count); - ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, token_ids_type, token_count); - ggml_tensor * output = ggml_new_tensor_2d(ctx, output_type, output_hidden_size, output_token_count); - REQUIRE(weight != nullptr); - REQUIRE(token_ids != nullptr); - REQUIRE(output != nullptr); + ggml_free(ctx); +} - ggml::hrx::Graph graph = build_manual_token_embedding_graph(weight, token_ids, output); - return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); +static void run_rmsnorm_binary_k16_output_checks() { + const struct { + ggml_type type; + int64_t hidden; + int64_t tokens; + int projections; + bool side_use; + bool observed; + bool inplace; + bool multiply; + 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 }, + { GGML_TYPE_Q6_K, 5120, 1024, 1, false, false, false, true, true }, + { GGML_TYPE_Q4_K, 5120, 512, 2, true, true, false, true, true }, + { GGML_TYPE_Q4_K, 6144, 512, 1, false, false, false, true, true }, + { GGML_TYPE_Q4_K, 7168, 512, 1, false, false, false, true, false }, + { GGML_TYPE_Q4_K, 4096, 256, 1, false, false, false, true, false }, + { GGML_TYPE_Q4_K, 5120, 5, 1, false, false, false, true, false }, + { GGML_TYPE_Q4_K, 4096, 1024, 1, false, false, false, true, false }, + { GGML_TYPE_Q4_K, 5120, 512, 0, true, true, false, true, false }, + { GGML_TYPE_F32, 5120, 512, 1, true, false, false, true, false }, + { GGML_TYPE_Q4_K, 5120, 512, 1, false, false, true, true, false }, + { GGML_TYPE_Q4_K, 5120, 512, 1, 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); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden); + ggml_tensor * rms = + test.inplace ? ggml_rms_norm_inplace(ctx, input, 1.0e-6f) : ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * norm = test.inplace ? ggml_mul_inplace(ctx, rms, weight) : + test.multiply ? ggml_mul(ctx, rms, weight) : + ggml_add(ctx, rms, weight); + if (test.observed) { + ggml_set_output(norm); + } + ggml_cgraph * graph = ggml_new_graph(ctx); + for (int i = 0; i < test.projections; ++i) { + ggml_tensor * projection = ggml_new_tensor_2d(ctx, test.type, test.hidden, 4096); + ggml_build_forward_expand(graph, ggml_mul_mat(ctx, projection, norm)); + } + ggml_tensor * side = test.side_use ? ggml_scale(ctx, norm, 0.5f) : nullptr; + if (side != nullptr) { + ggml_build_forward_expand(graph, side); + } + 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_binary_f32_k16"; + }); + REQUIRE((producer != plan.dispatches.end()) == test.packed); + const auto * output = imported.graph.values().find_tensor(norm); + REQUIRE(output != nullptr); + if (test.packed) { + const auto * generated = + plan.metadata.find_generated_resource(output->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + REQUIRE(generated != nullptr); + REQUIRE(generated->byte_count == output->byte_count / 2); + REQUIRE(producer->bindings.size() == 4); + REQUIRE(producer->bindings[2].value == output->id); + REQUIRE(producer->bindings[3].value == generated->generated_value); + REQUIRE(plan.metadata.find_alternate_value(output->id, GGML_TYPE_F16, output->byte_count / 2) == 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"; + }) == 0); + REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + 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); + REQUIRE(side_value != nullptr); + const auto consumer = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return !dispatch.bindings.empty() && dispatch.bindings.back().value == side_value->id; + }); + REQUIRE(consumer != plan.dispatches.end()); + REQUIRE(consumer->bindings.front().value == output->id); + } + 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 schedule_qwen_token_embedding_command(ggml_context * ctx, - ggml_tensor * output, - int64_t expected_token_count, - int64_t expected_vocabulary_count) { - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); +static void run_rmsnorm_binary_f16_output_checks() { + { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 1); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 7168, 512); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 7168); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * norm = ggml_mul(ctx, rms, norm_weight); + ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 7168, 2048); + ggml_tensor * output = ggml_mul_mat(ctx, projection_weight, norm); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); + 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_f32_f16"; + }); + REQUIRE(producer != plan.dispatches.end()); + REQUIRE(producer->bindings.size() == 4); + const ggml::hrx::Value * norm_value = imported.graph.values().find_tensor(norm); + REQUIRE(norm_value != nullptr); + const auto * f16 = + plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_F16, norm_value->byte_count / 2); + REQUIRE(f16 != nullptr); + REQUIRE(producer->bindings.back().value == f16->alternate_value); + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*producer && !dispatch.bindings.empty() && + dispatch.bindings.front().value == f16->alternate_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); + } - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == "qwen3_moe:qwen_token_embedding_q4k"); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); - REQUIRE(dispatch.kernel.integer_parameters.at("vocabulary_count") == expected_vocabulary_count); - REQUIRE(dispatch.bindings.size() == 3); + { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.front().bindings.size() == 3); - REQUIRE(commands.commands.front().bindings[0].name == "token_ids"); - REQUIRE(commands.commands.front().bindings[1].name == "weight"); - REQUIRE(commands.commands.front().bindings[2].name == "output"); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 512); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * norm = ggml_mul(ctx, rms, norm_weight); + ggml_tensor * q8_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, 2048, 1024); + ggml_tensor * f16_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 2048); + ggml_tensor * q8_output = ggml_mul_mat(ctx, q8_weight, norm); + ggml_tensor * f16_output = ggml_mul_mat(ctx, f16_weight, norm); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, q8_output); + ggml_build_forward_expand(graph, f16_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 auto * q8 = + plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(512, 2048)); + const auto * f16 = + plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_F16, norm_value->byte_count / 2); + REQUIRE(q8 != nullptr); + REQUIRE(f16 != nullptr); + REQUIRE(producer->bindings[3].value == q8->alternate_value); + REQUIRE(producer->bindings[4].value == f16->alternate_value); + REQUIRE(std::any_of(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*producer && !dispatch.bindings.empty() && + dispatch.bindings.front().value == q8->alternate_value; + })); + REQUIRE(std::any_of(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return &dispatch != &*producer && !dispatch.bindings.empty() && + dispatch.bindings.front().value == f16->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); + } + + { + 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_qwen_token_embedding_dispatch_checks() { +static void run_rmsnorm_gate_packed_output_checks() { ggml_init_params params = {}; - params.mem_size = 2 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); - ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - REQUIRE(weight != nullptr); - REQUIRE(token_ids != nullptr); - ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); - REQUIRE(output != nullptr); - schedule_qwen_token_embedding_command(ctx, output, 1, 151936); + const struct { + ggml_type type; + int64_t hidden; + int64_t channels; + int64_t tokens; + bool side_use; + bool packed; + int copies; + } cases[] = { + { GGML_TYPE_Q4_K, 128, 6144, 512, false, true, 0 }, + { GGML_TYPE_Q6_K, 256, 4096, 1024, false, true, 0 }, + { GGML_TYPE_Q6_K, 1024, 4096, 512, false, true, 0 }, + { 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) { + const int64_t heads = test.channels / test.hidden; + ggml_tensor * input = heads == 1 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens) : + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.hidden, heads, test.tokens); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden); + 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_tensor * reshaped = heads == 1 ? gated : ggml_reshape_2d(ctx, gated, test.channels, test.tokens); + ggml_tensor * projection = ggml_new_tensor_2d(ctx, test.type, test.channels, 4096); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_mul_mat(ctx, projection, reshaped)); + if (test.side_use) { + 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_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); + const auto * matmul_input = imported.graph.values().find_tensor(reshaped); + REQUIRE(output != nullptr && matmul_input != nullptr); + const auto * alternate = plan.metadata.find_alternate_value(output->id, GGML_TYPE_F16, output->byte_count / 2); + REQUIRE((alternate == nullptr) == test.packed); + if (test.packed) { + const auto * generated = + plan.metadata.find_generated_resource(matmul_input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + REQUIRE(generated != nullptr); + REQUIRE(producer->bindings.back().value == generated->generated_value); + REQUIRE(std::next(producer) != plan.dispatches.end()); + REQUIRE(std::next(producer)->bindings.front().value == generated->generated_value); + } + 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"; + }) == test.copies); + 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_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 151936); - ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 13); - REQUIRE(weight != nullptr); - REQUIRE(token_ids != nullptr); - ggml_tensor * output = ggml_get_rows(ctx, weight, token_ids); + 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); - schedule_qwen_token_embedding_command(ctx, output, 13, 151936); + 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); +} - REQUIRE( - !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_F32, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); - REQUIRE( - !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I64, GGML_TYPE_F32, 2048, 151936, 1)); - REQUIRE( - !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F16, 2048, 151936, 1)); - REQUIRE( - !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 1024, 151936, 1)); - REQUIRE( - !manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 3072, 248320, 1)); - REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q4_K, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, - 1, 2048, 2)); +static void run_rmsnorm_gate_q8_output_checks() { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + const struct { + ggml_type type; + int64_t hidden; + int64_t channels; + int64_t tokens; + int64_t outputs; + bool side_use; + bool terminal; + bool q8; + } cases[] = { + { GGML_TYPE_Q4_K, 128, 6144, 1, 5120, false, false, true }, + { GGML_TYPE_Q4_K, 128, 6144, 2, 5120, false, false, true }, + { GGML_TYPE_Q4_K, 128, 6144, 3, 5120, false, false, true }, + { GGML_TYPE_Q4_K, 128, 6144, 4, 5120, false, false, true }, + { GGML_TYPE_Q4_K, 128, 6144, 5, 5120, false, false, true }, + { GGML_TYPE_Q6_K, 256, 4096, 3, 256, false, false, true }, + { GGML_TYPE_Q4_K, 512, 1536, 5, 256, false, false, true }, + { GGML_TYPE_Q4_K, 1024, 1024, 4, 256, false, false, true }, + { GGML_TYPE_Q4_K, 256, 256, 1, 64, false, false, true }, + { GGML_TYPE_Q4_K, 128, 6144, 6, 5120, false, false, false }, + { GGML_TYPE_Q4_K, 128, 6144, 256, 5120, false, false, false }, + { GGML_TYPE_Q4_K, 128, 6144, 512, 5120, false, false, false }, + { GGML_TYPE_Q4_K, 128, 6144, 5, 5120, true, false, false }, + { GGML_TYPE_Q4_K, 128, 6144, 5, 48, false, false, false }, + { GGML_TYPE_Q6_K, 128, 6144, 5, 256, false, true, false }, + }; + + for (const auto & test : cases) { + const int64_t heads = test.channels / test.hidden; + ggml_tensor * input = heads == 1 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens) : + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.hidden, heads, test.tokens); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden); + 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_tensor * reshaped = heads == 1 ? gated : ggml_reshape_2d(ctx, gated, test.channels, test.tokens); + ggml_tensor * projection = ggml_new_tensor_2d(ctx, test.type, test.channels, test.outputs); + ggml_tensor * projected = ggml_mul_mat(ctx, projection, reshaped); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, test.terminal ? projected : ggml_scale(ctx, projected, 0.5f)); + if (test.side_use) { + ggml_build_forward_expand(graph, ggml_scale(ctx, gated, 0.25f)); + } + 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) { + 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); + const auto * alternate = plan.metadata.find_alternate_value(output->id, GGML_TYPE_Q8_1, bytes); + REQUIRE(alternate != nullptr); + REQUIRE(producer->bindings.back().value == alternate->alternate_value); + REQUIRE(producer->bindings.back().length == bytes); + 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 std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), + [&](const auto & binding) { return binding.value == alternate->alternate_value; }); + })); + } + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } ggml_free(ctx); } -static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, - int64_t hidden_size, - int64_t source_token_count, - int64_t output_token_count, - bool shared_row_ids, - int64_t second_source_hidden_size = -1) { - if (second_source_hidden_size < 0) { - second_source_hidden_size = hidden_size; - } +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); - ggml::hrx::Graph graph; + 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)); + } - ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); - ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, second_source_hidden_size, source_token_count); - ggml_tensor * row_ids0 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); - ggml_tensor * row_ids1 = shared_row_ids ? row_ids0 : ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); - ggml_tensor * selected0 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); - ggml_tensor * selected1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); - ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, output_token_count); - REQUIRE(attention != nullptr); - REQUIRE(residual != nullptr); - REQUIRE(row_ids0 != nullptr); - REQUIRE(row_ids1 != nullptr); - REQUIRE(selected0 != nullptr); - REQUIRE(selected1 != nullptr); - REQUIRE(output != nullptr); + 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); + } +} - const ggml::hrx::ValueId attention_value = - graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId residual_value = - graph.values().get_or_add_tensor_value(residual, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId row_ids0_value = - graph.values().get_or_add_tensor_value(row_ids0, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId row_ids1_value = - graph.values().get_or_add_tensor_value(row_ids1, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId selected0_value = - graph.values().get_or_add_tensor_value(selected0, ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId selected1_value = - graph.values().get_or_add_tensor_value(selected1, ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId output_value = - graph.values().get_or_add_tensor_value(output, ggml::hrx::ValueKind::External); +static void run_quantized_conv4_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); - graph.add_node(GGML_OP_GET_ROWS, selected0_value, { attention_value, row_ids0_value }); - graph.add_node(GGML_OP_GET_ROWS, selected1_value, { residual_value, row_ids1_value }); - graph.add_node(GGML_OP_ADD, output_value, { selected0_value, selected1_value }); - REQUIRE(graph.build_index().success()); - return graph; -} + const struct { + ggml_type type; + int64_t inputs; + int64_t channels; + int64_t tokens; + bool side_use; + bool exposed; + bool fused; + bool gathered_state = false; + } cases[] = { + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, false, true }, + { GGML_TYPE_Q6_K, 5120, 10240, 512, false, false, true }, + { GGML_TYPE_Q4_K, 4096, 8192, 512, false, false, true }, + { GGML_TYPE_Q6_K, 4096, 8192, 512, false, false, true }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false }, + { GGML_TYPE_Q4_K, 5120, 10240, 256, false, false, false }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, false, true, true }, + { GGML_TYPE_Q6_K, 5120, 10240, 512, false, false, true, true }, + { GGML_TYPE_Q4_K, 4096, 8192, 512, false, false, true, true }, + { GGML_TYPE_Q6_K, 4096, 8192, 512, false, false, true, true }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false, true }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false, true }, + }; -static bool manual_gather_add_graph_is_supported(ggml_context * ctx, - int64_t hidden_size, - int64_t source_token_count, - int64_t output_token_count, - bool shared_row_ids, - int64_t second_source_hidden_size = -1) { - ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, hidden_size, source_token_count, output_token_count, - shared_row_ids, second_source_hidden_size); - return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); + for (const auto & test : cases) { + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, test.type, test.inputs, test.channels); + ggml_tensor * raw = ggml_mul_mat(ctx, weight, input); + if (test.exposed) { + ggml_set_output(raw); + } + ggml_tensor * x = ggml_reshape_3d(ctx, raw, test.channels, test.tokens, 1); + ggml_tensor * state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3, test.channels, 1); + if (test.gathered_state) { + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + state = ggml_get_rows(ctx, ggml_reshape_2d(ctx, state, 3 * test.channels, 1), indices); + state = ggml_reshape_3d(ctx, state, 3, test.channels, 1); + } + ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, test.channels); + ggml_tensor * window = ggml_concat(ctx, state, ggml_transpose(ctx, x), 0); + ggml_tensor * tail = + ggml_view_3d(ctx, window, 3, test.channels, 1, window->nb[1], window->nb[2], test.tokens * sizeof(float)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3 * test.channels, 1); + ggml_tensor * target = ggml_view_2d(ctx, cache, 3 * test.channels, 1, cache->nb[1], 0); + ggml_tensor * output = ggml_silu(ctx, ggml_ssm_conv(ctx, window, filter)); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_cpy(ctx, tail, target)); + ggml_build_forward_expand(graph, output); + if (test.side_use) { + ggml_build_forward_expand(graph, ggml_scale(ctx, raw, 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 fused = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + (test.gathered_state ? "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" : + "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4"); + }); + REQUIRE((fused != plan.dispatches.end()) == test.fused); + if (test.fused) { + REQUIRE(fused->bindings.size() == (test.gathered_state ? 5 : 6)); + require_compile_parameter(*fused, "ggml.mul_mat.weight_format", test.type == GGML_TYPE_Q4_K ? "4" : "6"); + if (test.gathered_state) { + const auto finish = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32"; + }); + const auto gather = + std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id).find("get_rows") != std::string::npos; + }); + REQUIRE(finish != plan.dispatches.end()); + REQUIRE(gather != plan.dispatches.end()); + REQUIRE(fused < gather && gather < finish); + REQUIRE(finish->bindings[2].value == fused->bindings[4].value); + REQUIRE(finish->bindings[3].value == fused->bindings[3].value); + REQUIRE(finish->bindings[2].length == static_cast(6 * test.channels) * sizeof(float)); + } + REQUIRE(std::none_of(plan.transients.begin(), plan.transients.end(), [](const auto & transient) { + return transient.name == "llm.ssm_conv.window_snapshot"; + })); + } + 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 bool partial_gather_add_graph_is_supported(ggml_context * ctx) { - ggml::hrx::Graph graph; +static void run_generic_ssm_conv_binary_dispatch_checks() { + const struct { + ggml::hrx::BinaryKind kind; + bool conv_is_lhs; + } cases[] = { + { ggml::hrx::BinaryKind::Add, true }, + { ggml::hrx::BinaryKind::Sub, true }, + { ggml::hrx::BinaryKind::Sub, false }, + { ggml::hrx::BinaryKind::Mul, true }, + { ggml::hrx::BinaryKind::Div, true }, + { ggml::hrx::BinaryKind::Div, false }, + }; - ggml_tensor * attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 13); - ggml_tensor * row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * add_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * add_output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - REQUIRE(attention != nullptr); - REQUIRE(row_ids != nullptr); - REQUIRE(selected != nullptr); - REQUIRE(add_input != nullptr); - REQUIRE(add_output != nullptr); + 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); - const ggml::hrx::ValueId attention_value = - graph.values().get_or_add_tensor_value(attention, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId row_ids_value = - graph.values().get_or_add_tensor_value(row_ids, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId selected_value = - graph.values().get_or_add_tensor_value(selected, ggml::hrx::ValueKind::Transient); - const ggml::hrx::ValueId add_input_value = - graph.values().get_or_add_tensor_value(add_input, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId add_output_value = - graph.values().get_or_add_tensor_value(add_output, ggml::hrx::ValueKind::External); + constexpr int64_t d_conv = 4; + constexpr int64_t d_inner = 64; + constexpr int64_t n_t = 8; + ggml_tensor * window = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_conv + n_t - 1, d_inner); + ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_conv, d_inner); + ggml_tensor * operand = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_inner, n_t); + ggml_tensor * conv = ggml_ssm_conv(ctx, window, filter); + ggml_tensor * output = nullptr; + switch (test.kind) { + case ggml::hrx::BinaryKind::Add: + output = ggml_add(ctx, conv, operand); + break; + case ggml::hrx::BinaryKind::Sub: + output = test.conv_is_lhs ? ggml_sub(ctx, conv, operand) : ggml_sub(ctx, operand, conv); + break; + case ggml::hrx::BinaryKind::Mul: + output = ggml_mul(ctx, conv, operand); + break; + case ggml::hrx::BinaryKind::Div: + output = test.conv_is_lhs ? ggml_div(ctx, conv, operand) : ggml_div(ctx, operand, conv); + break; + default: + REQUIRE(false); + } - graph.add_node(GGML_OP_GET_ROWS, selected_value, { attention_value, row_ids_value }); - graph.add_node(GGML_OP_ADD, add_output_value, { selected_value, add_input_value }); - REQUIRE(graph.build_index().success()); - return ggml::hrx::DispatchScheduler::can_schedule_graph(graph, test_dispatch_target()); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + 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))); + require_compile_parameter(dispatch, "llm.ssm_conv.generic.binary_lhs", test.conv_is_lhs ? "1" : "0"); + REQUIRE(dispatch.bindings[2].value == imported.graph.values().find_tensor(operand)->id); + REQUIRE(dispatch.bindings[3].value == imported.graph.values().find_tensor(output)->id); + + 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 schedule_gather_add_command(ggml::hrx::Graph & graph, - int64_t expected_hidden_size, - int64_t expected_source_token_count, - int64_t expected_output_token_count) { - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_gather_add_f32"); - REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == expected_hidden_size); - REQUIRE(dispatch.kernel.integer_parameters.at("source_token_count") == expected_source_token_count); - REQUIRE(dispatch.kernel.integer_parameters.at("output_token_count") == expected_output_token_count); - REQUIRE(dispatch.bindings.size() == 4); +static void run_lfm_dconv3_dispatch_checks() { + const struct { + int64_t tokens; + int64_t hidden; + int64_t state_rows; + int64_t filter_rows; + bool disabled; + bool fused; + } cases[] = { + {1, 2048, 2, 3, false, true}, + {64, 2048, 2, 3, false, true}, + {2, 2048, 2, 3, false, false}, + {64, 2016, 2, 3, false, false}, + {64, 2048, 3, 3, false, false}, + {64, 2048, 2, 4, false, false}, + {64, 2048, 2, 3, true, false}, + }; - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.front().bindings.size() == 4); - REQUIRE(commands.commands.front().bindings[0].name == "attention"); - REQUIRE(commands.commands.front().bindings[1].name == "residual"); - REQUIRE(commands.commands.front().bindings[2].name == "output_ids"); - REQUIRE(commands.commands.front().bindings[3].name == "output"); + 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 * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.hidden, test.tokens, 1); + ggml_tensor * state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.state_rows, test.hidden, 1); + ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.filter_rows, test.hidden); + ggml_tensor * window = ggml_concat(ctx, state, ggml_transpose(ctx, x), 0); + ggml_tensor * tail = ggml_view_3d(ctx, window, test.state_rows, test.hidden, 1, window->nb[1], window->nb[2], + test.tokens * sizeof(float)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.state_rows * test.hidden, 1); + ggml_tensor * target = ggml_view_2d(ctx, cache, test.state_rows * test.hidden, 1, cache->nb[1], 0); + ggml_tensor * output = ggml_ssm_conv(ctx, window, filter); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_cpy(ctx, tail, target)); + ggml_build_forward_expand(graph, output); + if (test.disabled) { + REQUIRE(setenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION", "1", 1) == 0); + } + 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 fused = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_lfm_dconv3_state_f32"; + }); + REQUIRE((fused != plan.dispatches.end()) == test.fused); + if (test.fused) { + REQUIRE(fused->bindings.size() == 5); + require_compile_parameter(*fused, "llm.ssm_conv.lfm_dconv3.n_t", std::to_string(test.tokens)); + } + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + if (test.disabled) { + REQUIRE(unsetenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION") == 0); + } + ggml_free(ctx); + } } -static void run_gather_add_dispatch_checks() { +static void run_gdn_native_projection_pair_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - { - ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, 2048, 13, 1, true); - schedule_gather_add_command(graph, 2048, 13, 1); - } - { - ggml::hrx::Graph graph = build_manual_gather_add_graph(ctx, 2048, 128, 8, true); - schedule_gather_add_command(graph, 2048, 128, 8); - } - - REQUIRE(!manual_gather_add_graph_is_supported(ctx, 2048, 13, 1, false)); - REQUIRE(!manual_gather_add_graph_is_supported(ctx, 2048, 13, 1, true, 1024)); - REQUIRE(!manual_gather_add_graph_is_supported(ctx, 96, 13, 1, true)); - REQUIRE(!partial_gather_add_graph_is_supported(ctx)); + const struct { + int64_t tokens; + int64_t inputs; + int64_t heads; + bool different_input; + bool mutate_input; + bool aliased_weight; + ggml_type weight_type; + bool fused; + } cases[] = { + {1, 5120, 48, false, false, false, GGML_TYPE_Q4_K, true}, + {1, 4096, 24, false, false, false, GGML_TYPE_Q4_K, true}, + {1, 8192, 48, false, false, false, GGML_TYPE_Q4_K, true}, + {1, 32768, 48, false, false, false, GGML_TYPE_Q4_K, true}, + {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}, + {1, 5120, 48, true, false, false, GGML_TYPE_Q4_K, false}, + {1, 5120, 48, false, true, false, GGML_TYPE_Q4_K, false}, + {1, 5120, 48, false, false, true, GGML_TYPE_Q4_K, false}, + {1, 5120, 48, false, false, false, GGML_TYPE_Q6_K, false}, + }; + for (const auto & test : cases) { + const int64_t width = 128; + const int64_t qheads = test.heads / 4; + const int64_t hidden = width * (2 * qheads + test.heads); + const int64_t state_elements = width * width * test.heads; + const size_t state_bytes = state_elements * sizeof(float); + ggml_tensor * input = ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens), 0.5f); + ggml_tensor * first_weight = ggml_new_tensor_2d(ctx, test.weight_type, test.inputs, test.heads); + ggml_tensor * second_weight = ggml_dup_tensor(ctx, first_weight); + if (test.aliased_weight) { + second_weight = ggml_view_2d(ctx, second_weight, test.inputs, test.heads, second_weight->nb[1], 0); + } + ggml_tensor * alpha = ggml_mul_mat(ctx, first_weight, input); + ggml_tensor * second_input = test.different_input ? ggml_dup_tensor(ctx, input) : input; + ggml_tensor * beta_raw = ggml_mul_mat(ctx, second_weight, second_input); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.heads); + ggml_tensor * scale = ggml_dup_tensor(ctx, bias); + ggml_tensor * gate = ggml_mul( + ctx, ggml_softplus(ctx, ggml_add(ctx, ggml_reshape_3d(ctx, alpha, test.heads, test.tokens, 1), bias)), + scale); + gate = ggml_reshape_4d(ctx, gate, 1, test.heads, test.tokens, 1); + ggml_tensor * beta = ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta_raw, 1, test.heads, test.tokens, 1)); + ggml_tensor * raw = ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, test.tokens), 0.5f); + const auto view = [&](int64_t heads, size_t offset) { + return ggml_view_4d(ctx, raw, width, heads, test.tokens, 1, width * sizeof(float), hidden * sizeof(float), + hidden * test.tokens * sizeof(float), offset); + }; + ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); + ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); + ggml_tensor * v = view(test.heads, 2 * width * qheads * sizeof(float)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 20); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * gathered = ggml_get_rows(ctx, cache, ids); + ggml_tensor * state = ggml_reshape_4d(ctx, gathered, width, width, test.heads, 1); + ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, gate, beta, state, 5); + const int64_t count = std::min(test.tokens, 5); + const size_t attention_bytes = width * test.heads * test.tokens * sizeof(float); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, test.heads, test.tokens, 1, width * sizeof(float), + width * test.heads * sizeof(float), attention_bytes, 0); + ggml_tensor * snapshots = + ggml_view_3d(ctx, gdn, state_elements, 1, count, state_bytes, state_bytes, attention_bytes); + ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, 1, count, state_bytes, 4 * state_bytes, 0); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, raw); + ggml_build_forward_expand(graph, alpha); + if (test.mutate_input) { + ggml_build_forward_expand(graph, ggml_cpy(ctx, ggml_dup_tensor(ctx, input), input)); + } + ggml_build_forward_expand(graph, beta_raw); + ggml_build_forward_expand(graph, gathered); + ggml_build_forward_expand(graph, ggml_cpy(ctx, snapshots, target)); + ggml_build_forward_expand(graph, ggml_scale(ctx, attention, 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 fused = 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_dual_q4_f32_decode"; + }); + REQUIRE((fused != plan.dispatches.end()) == test.fused); + if (test.fused) { + REQUIRE(fused->bindings.size() == 5); + require_compile_parameter(*fused, "ggml.mul_mat_dual_q4_f32_c1.input_size", std::to_string(test.inputs)); + REQUIRE(fused->bindings[0].value == imported.graph.values().find_tensor(input)->id); + REQUIRE(fused->bindings[1].value == imported.graph.values().find_tensor(first_weight)->id); + REQUIRE(fused->bindings[2].value == imported.graph.values().find_tensor(second_weight)->id); + 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()); + REQUIRE(command_program_verifies(commands)); + } ggml_free(ctx); } -static void schedule_qwen_flash_attention_command(ggml_context * ctx, ggml_tensor * output) { - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); +static void run_gdn_selected_snapshot_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 1); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_FLASH_ATTN_EXT); + const struct { + int64_t tokens; + int64_t sequences; + bool ready; + bool side_use; + bool cache_read; + bool observed; + bool aligned; + bool fused; + } cases[] = { + {2,1,true,false,false,false,true,true}, + {3,1,true,false,false,false,true,true}, + {4,1,true,false,false,false,true,true}, + {5,1,true,false,false,false,true,true}, + {1,1,true,false,false,false,true,false}, + {5,2,true,false,false,false,true,false}, + {5,1,false,false,false,false,true,false}, + {5,1,true,true,false,false,true,false}, + {5,1,true,false,true,false,true,false}, + {5,1,true,false,false,true,true,false}, + {5,1,true,false,false,false,false,false}, + }; + for (const auto & test : cases) { + const int64_t width=128, heads=6, qheads=2, hidden=1280; + const int64_t state_elements=width*width*heads; + const size_t state_bytes=state_elements*sizeof(float); + ggml_tensor * raw=ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, test.tokens*test.sequences), 0.5f); + ggml_tensor * alpha=ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, heads, test.tokens*test.sequences), 0.5f); + ggml_tensor * beta_raw=ggml_scale(ctx, ggml_dup_tensor(ctx, alpha), 0.5f); + ggml_tensor * bias=ggml_new_tensor_1d(ctx, GGML_TYPE_F32, heads); + ggml_tensor * scale=ggml_dup_tensor(ctx, bias); + ggml_tensor * gate=ggml_mul(ctx, ggml_softplus(ctx, ggml_add(ctx, ggml_reshape_3d(ctx, alpha, heads, test.tokens, test.sequences), bias)), scale); + gate=ggml_reshape_4d(ctx, gate, 1, heads, test.tokens, test.sequences); + ggml_tensor * beta=ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta_raw, 1, heads, test.tokens, test.sequences)); + ggml_tensor * cache=ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 20); + ggml_tensor * ids=ggml_new_tensor_1d(ctx, GGML_TYPE_I32, test.sequences); + ggml_tensor * gathered=ggml_get_rows(ctx, cache, ids); + if (test.observed) { + ggml_set_output(gathered); + } + ggml_tensor * state=ggml_reshape_4d(ctx, gathered, width, width, heads, test.sequences); + const auto view=[&](int64_t nheads, size_t offset) { + return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, + width*sizeof(float), hidden*sizeof(float), hidden*test.tokens*sizeof(float), offset); + }; + ggml_tensor * q=ggml_l2_norm(ctx, view(qheads,0), 1.e-6f); + ggml_tensor * k=ggml_l2_norm(ctx, view(qheads,width*qheads*sizeof(float)), 1.e-6f); + ggml_tensor * v=view(heads,2*width*qheads*sizeof(float)); + ggml_tensor * gdn=ggml_gated_delta_net(ctx,q,k,v,gate,beta,state,5); + const size_t attention_bytes=width*heads*test.tokens*test.sequences*sizeof(float); + ggml_tensor * attention=ggml_view_4d(ctx,gdn,width,heads,test.tokens,test.sequences, + width*sizeof(float),width*heads*sizeof(float), + width*heads*test.tokens*sizeof(float),0); + ggml_tensor * snapshots=ggml_view_3d(ctx,gdn,state_elements,test.sequences,test.tokens, + state_bytes,state_bytes*test.sequences,attention_bytes); + ggml_tensor * target=ggml_view_3d(ctx,cache,state_elements,test.sequences,test.tokens, + state_bytes,4*state_bytes,test.aligned ? state_bytes : sizeof(float)); + ggml_cgraph * graph=ggml_new_graph(ctx); + if (test.ready) { + ggml_build_forward_expand(graph,raw); + } + ggml_build_forward_expand(graph,alpha); + ggml_build_forward_expand(graph,beta_raw); + ggml_build_forward_expand(graph,gathered); + if (test.cache_read) { + ggml_build_forward_expand(graph,ggml_scale(ctx,cache,0.25f)); + } + ggml_build_forward_expand(graph,ggml_cpy(ctx,snapshots,target)); + ggml_build_forward_expand(graph,ggml_scale(ctx,attention,0.5f)); + if (test.side_use) { + ggml_build_forward_expand(graph,ggml_scale(ctx,gathered,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 fused=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_f32_wmma_head128_selected_snapshot_projection_epilogue"; + }); + REQUIRE((fused!=plan.dispatches.end())==test.fused); + if (test.fused) { + REQUIRE(fused->bindings.size()==11); + require_compile_parameter(*fused,"llm.gated_delta_net.state_row_count","20"); + REQUIRE(fused->bindings[7].value==imported.graph.values().find_tensor(cache)->id); + REQUIRE(fused->bindings[10].value==imported.graph.values().find_tensor(ids)->id); + REQUIRE(std::none_of(plan.dispatches.begin(),plan.dispatches.end(),[](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id)=="loom_libs:ggml_get_rows_f32"; + })); + } + 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); +} - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); +static void run_gdn_rmsnorm_gate_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); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == "qwen3_moe:qwen3_moe_flash_attention_f32_f16_wmma"); - REQUIRE(dispatch.kernel.integer_parameters.at("query_token_count") == 4); - REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == 8); - REQUIRE(dispatch.bindings.size() == 5); - require_compile_parameter(dispatch, "qwen3_moe.attention.query_head_count", "4"); - require_compile_parameter(dispatch, "qwen3_moe.attention.key_value_head_count", "2"); - require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", "4"); + const struct { + ggml_type type; + int64_t tokens; + bool side_use; + bool attention_output; + bool gdn_output; + bool fused; + } cases[] = { + { GGML_TYPE_Q4_K, 512, false, false, false, true }, + { GGML_TYPE_Q6_K, 512, false, false, false, true }, + { GGML_TYPE_Q4_K, 512, true, false, false, false }, + { GGML_TYPE_Q4_K, 256, false, false, false, false }, + { GGML_TYPE_Q4_K, 512, false, true, false, false }, + { GGML_TYPE_Q4_K, 512, false, false, true, false }, + }; - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.front().bindings.size() == 5); - REQUIRE(commands.commands.front().bindings[0].name == "query"); - 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 == "output"); + for (const auto & test : cases) { + const int64_t width = 128; + const int64_t heads = 2; + const int64_t channels = width * heads; + const int64_t state_elements = width * channels; + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, 1, test.tokens, 1); + ggml_tensor * k = ggml_dup_tensor(ctx, q); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, heads, test.tokens, 1); + ggml_tensor * g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, heads, test.tokens, 1); + ggml_tensor * beta = ggml_dup_tensor(ctx, g); + ggml_tensor * state = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, width, heads, 1); + ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, 1); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, 1, width * sizeof(float), + channels * sizeof(float), channels * test.tokens * sizeof(float), 0); + if (test.attention_output) { + ggml_set_output(attention); + } + if (test.gdn_output) { + ggml_set_output(gdn); + } + ggml_tensor * new_state = ggml_view_2d(ctx, gdn, state_elements, 1, state_elements * sizeof(float), + channels * test.tokens * sizeof(float)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); + ggml_tensor * gate = ggml_dup_tensor(ctx, attention); + ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, attention, 2.0e-6f), weight); + ggml_tensor * gated = ggml_mul(ctx, norm, ggml_silu(ctx, gate)); + ggml_tensor * input = ggml_reshape_2d(ctx, gated, channels, test.tokens); + ggml_tensor * projection = ggml_new_tensor_2d(ctx, test.type, channels, 4096); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_cpy(ctx, new_state, cache)); + ggml_build_forward_expand(graph, ggml_mul_mat(ctx, projection, input)); + if (test.side_use) { + ggml_build_forward_expand(graph, ggml_scale(ctx, attention, 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 fused = 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_f32_wmma_head128_rmsnorm_gate"; + }); + 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_publish"; + }) == (test.fused ? 0 : 1)); + if (test.fused) { + REQUIRE(fused->bindings.size() == 11); + const auto * projected = imported.graph.values().find_tensor(input); + REQUIRE(projected != nullptr); + const auto * packed = + plan.metadata.find_generated_resource(projected->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + REQUIRE(packed != nullptr && fused->bindings.back().value == packed->generated_value); + require_compile_parameter(*fused, "ggml.rmsnorm_gate_f32.rms_epsilon", "1.99999999e-06"); + } + 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_qwen_flash_attention_dispatch_checks() { +static void run_qwen_matmul_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; + params.mem_size = 12 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2); - schedule_qwen_flash_attention_command(ctx, output); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } + { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 8, 4, 2); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 2048, 128); } + { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F32); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, + 128); } + { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, false); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 3); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 2048, + 128); } + { - ggml_tensor * output = - build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, true); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 5); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 2048, + 128); } + { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, - false, 64, 1.0f / std::sqrt(64.0f)); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 6); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 6, 2048, 128); } + { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, - false, kQwenFlashHeadSize, 1.0f); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, 640, 128); } + { - ggml_tensor * output = - build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, true, false, - kQwenFlashHeadSize, 1.0f / std::sqrt(128.0f), 1.0f); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, 2048, 128); } - ggml_free(ctx); -} - -struct QwenAttentionPostprocessTensors { - ggml_tensor * query_raw = nullptr; - ggml_tensor * key_raw = nullptr; - ggml_tensor * value_raw = nullptr; - ggml_tensor * query_reshape = nullptr; - ggml_tensor * key_reshape = nullptr; - ggml_tensor * value_reshape = nullptr; - ggml_tensor * query_output = nullptr; - ggml_tensor * key_cache = nullptr; - ggml_tensor * value_cache = nullptr; - ggml_tensor * key_output = nullptr; - ggml_tensor * value_output = nullptr; - ggml_tensor * positions = nullptr; - ggml_tensor * key_cache_indices = nullptr; - ggml_tensor * value_cache_indices = nullptr; - ggml_tensor * mask = nullptr; - ggml_tensor * flash_output = nullptr; -}; - -static QwenAttentionPostprocessTensors build_qwen_attention_postprocess_graph(ggml_context * ctx, - int64_t token_count, - int64_t query_head_count, - int64_t key_value_head_count, - int64_t cache_row_count, - float rms_epsilon = 0.000001f, - bool include_inverse_frequencies = true) { - QwenAttentionPostprocessTensors tensors; - const int64_t query_size = query_head_count * kQwenFlashHeadSize; - const int64_t key_value_size = key_value_head_count * kQwenFlashHeadSize; - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, token_count); - ggml_tensor * query_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, query_size); - ggml_tensor * key_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, key_value_size); - ggml_tensor * value_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, kQwenMoeHiddenSize, key_value_size); - REQUIRE(input != nullptr); - REQUIRE(query_weight != nullptr); - REQUIRE(key_weight != nullptr); - REQUIRE(value_weight != nullptr); - - tensors.query_raw = ggml_mul_mat(ctx, query_weight, input); - tensors.key_raw = ggml_mul_mat(ctx, key_weight, input); - tensors.value_raw = ggml_mul_mat(ctx, value_weight, input); - REQUIRE(tensors.query_raw != nullptr); - REQUIRE(tensors.key_raw != nullptr); - REQUIRE(tensors.value_raw != nullptr); - - tensors.query_reshape = ggml_reshape_3d(ctx, tensors.query_raw, kQwenFlashHeadSize, query_head_count, token_count); - tensors.key_reshape = ggml_reshape_3d(ctx, tensors.key_raw, kQwenFlashHeadSize, key_value_head_count, token_count); - tensors.value_reshape = - ggml_reshape_3d(ctx, tensors.value_raw, kQwenFlashHeadSize, key_value_head_count, token_count); - REQUIRE(tensors.query_reshape != nullptr); - REQUIRE(tensors.key_reshape != nullptr); - REQUIRE(tensors.value_reshape != nullptr); - - ggml_tensor * query_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); - ggml_tensor * key_norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize); - tensors.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * inverse_frequencies = - include_inverse_frequencies ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize / 2) : nullptr; - REQUIRE(query_norm_weight != nullptr); - REQUIRE(key_norm_weight != nullptr); - REQUIRE(tensors.positions != nullptr); - REQUIRE(include_inverse_frequencies == (inverse_frequencies != nullptr)); - - ggml_tensor * query_norm = ggml_rms_norm(ctx, tensors.query_reshape, rms_epsilon); - ggml_tensor * query_mul = ggml_mul(ctx, query_norm, query_norm_weight); - REQUIRE(query_norm != nullptr); - REQUIRE(query_mul != nullptr); - tensors.query_output = include_inverse_frequencies ? - ggml_rope_ext(ctx, query_mul, tensors.positions, inverse_frequencies, kQwenFlashHeadSize, - GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f) : - ggml_rope(ctx, query_mul, tensors.positions, kQwenFlashHeadSize, GGML_ROPE_TYPE_NEOX); - REQUIRE(tensors.query_output != nullptr); - - ggml_tensor * key_norm = ggml_rms_norm(ctx, tensors.key_reshape, rms_epsilon); - ggml_tensor * key_mul = ggml_mul(ctx, key_norm, key_norm_weight); - REQUIRE(key_norm != nullptr); - REQUIRE(key_mul != nullptr); - ggml_tensor * key_rope = include_inverse_frequencies ? - ggml_rope_ext(ctx, key_mul, tensors.positions, inverse_frequencies, kQwenFlashHeadSize, - GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f) : - ggml_rope(ctx, key_mul, tensors.positions, kQwenFlashHeadSize, GGML_ROPE_TYPE_NEOX); - REQUIRE(key_rope != nullptr); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + } - ggml_tensor * key_cache_rows = - ggml_reshape_2d(ctx, key_rope, kQwenFlashHeadSize * key_value_head_count, token_count); - ggml_tensor * value_cache_rows = - ggml_reshape_2d(ctx, tensors.value_reshape, kQwenFlashHeadSize * key_value_head_count, token_count); - REQUIRE(key_cache_rows != nullptr); - REQUIRE(value_cache_rows != nullptr); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); + } - tensors.key_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); - tensors.value_cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_size, cache_row_count); - tensors.key_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); - tensors.value_cache_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); - REQUIRE(tensors.key_cache != nullptr); - REQUIRE(tensors.value_cache != nullptr); - REQUIRE(tensors.key_cache_indices != nullptr); - REQUIRE(tensors.value_cache_indices != nullptr); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + } - tensors.key_output = ggml_set_rows(ctx, tensors.key_cache, key_cache_rows, tensors.key_cache_indices); - tensors.value_output = ggml_set_rows(ctx, tensors.value_cache, value_cache_rows, tensors.value_cache_indices); - REQUIRE(tensors.key_output != nullptr); - REQUIRE(tensors.value_output != nullptr); - return tensors; -} + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); + } -static ggml::hrx::GraphImportResult import_qwen_attention_postprocess_graph( - ggml_context * ctx, - const QwenAttentionPostprocessTensors & tensors) { - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, tensors.query_output); - ggml_build_forward_expand(graph, tensors.key_output); - ggml_build_forward_expand(graph, tensors.value_output); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + } - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - return imported; -} + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); + } -static ggml_tensor * append_qwen_flash_attention_consumer(ggml_context * ctx, - QwenAttentionPostprocessTensors & tensors, - int64_t token_count, - int64_t query_head_count, - int64_t key_value_head_count, - int64_t cache_row_count) { - ggml_tensor * query_layout = - ggml_reshape_3d(ctx, tensors.query_output, kQwenFlashHeadSize, query_head_count, token_count); - ggml_tensor * query_permute = ggml_permute(ctx, query_layout, 0, 2, 1, 3); - ggml_tensor * key_cache_layout = - ggml_reshape_3d(ctx, tensors.key_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); - ggml_tensor * key_permute = ggml_permute(ctx, key_cache_layout, 0, 2, 1, 3); - ggml_tensor * value_cache_layout = - ggml_reshape_3d(ctx, tensors.value_cache, kQwenFlashHeadSize, key_value_head_count, cache_row_count); - ggml_tensor * value_permute = ggml_permute(ctx, value_cache_layout, 0, 2, 1, 3); - tensors.mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, cache_row_count, token_count); - REQUIRE(query_layout != nullptr); - REQUIRE(query_permute != nullptr); - REQUIRE(key_cache_layout != nullptr); - REQUIRE(key_permute != nullptr); - REQUIRE(value_cache_layout != nullptr); - REQUIRE(value_permute != nullptr); - REQUIRE(tensors.mask != nullptr); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + } - tensors.flash_output = ggml_flash_attn_ext(ctx, query_permute, key_permute, value_permute, tensors.mask, - 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)), 0.0f, 0.0f); - REQUIRE(tensors.flash_output != nullptr); - return tensors.flash_output; -} + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); + } -static void schedule_qwen_attention_postprocess_command(ggml_context * ctx, - const QwenAttentionPostprocessTensors & tensors, - int64_t token_count, - int64_t query_head_count, - int64_t key_value_head_count, - int64_t cache_row_count, - bool expect_synthetic_inverse_frequencies = false) { - ggml::hrx::GraphImportResult imported = import_qwen_attention_postprocess_graph(ctx, tensors); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + } - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().initialization_dispatches.empty()); - REQUIRE(scheduler.plan().dispatches.size() == 4); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); + } - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == "qwen3_moe:qwen3_moe_attention_postprocess_f32_f16"); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); - REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_row_count); - REQUIRE(dispatch.bindings.size() == 12); - require_compile_parameter(dispatch, "qwen3_moe.model.rms_epsilon", "0.000001"); - require_compile_parameter(dispatch, "qwen3_moe.attention.head_size", std::to_string(kQwenFlashHeadSize)); - require_compile_parameter(dispatch, "qwen3_moe.attention.query_size", - std::to_string(query_head_count * kQwenFlashHeadSize)); - require_compile_parameter(dispatch, "qwen3_moe.attention.key_value_size", - std::to_string(key_value_head_count * kQwenFlashHeadSize)); - require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 256); + } - const ggml::hrx::Value * positions_value = imported.graph.values().find_tensor(tensors.positions); - const ggml::hrx::Value * key_indices_value = imported.graph.values().find_tensor(tensors.key_cache_indices); - const ggml::hrx::Value * value_indices_value = imported.graph.values().find_tensor(tensors.value_cache_indices); - const ggml::hrx::Value * query_raw_value = imported.graph.values().find_tensor(tensors.query_raw); - const ggml::hrx::Value * key_raw_value = imported.graph.values().find_tensor(tensors.key_raw); - const ggml::hrx::Value * value_raw_value = imported.graph.values().find_tensor(tensors.value_raw); - const ggml::hrx::Value * query_output_value = imported.graph.values().find_tensor(tensors.query_output); - const ggml::hrx::Value * key_cache_value = imported.graph.values().find_tensor(tensors.key_cache); - const ggml::hrx::Value * value_cache_value = imported.graph.values().find_tensor(tensors.value_cache); - REQUIRE(positions_value != nullptr); - REQUIRE(key_indices_value != nullptr); - REQUIRE(value_indices_value != nullptr); - REQUIRE(query_raw_value != nullptr); - REQUIRE(key_raw_value != nullptr); - REQUIRE(value_raw_value != nullptr); - REQUIRE(query_output_value != nullptr); - REQUIRE(key_cache_value != nullptr); - REQUIRE(value_cache_value != nullptr); - REQUIRE(dispatch.bindings[0].value == positions_value->id); - REQUIRE(dispatch.bindings[1].value == key_indices_value->id); - REQUIRE(dispatch.bindings[2].value == value_indices_value->id); - REQUIRE(dispatch.bindings[3].value == query_raw_value->id); - REQUIRE(dispatch.bindings[4].value == key_raw_value->id); - REQUIRE(dispatch.bindings[5].value == value_raw_value->id); - REQUIRE(dispatch.bindings[9].value == query_output_value->id); - REQUIRE(dispatch.bindings[10].value == key_cache_value->id); - REQUIRE(dispatch.bindings[11].value == value_cache_value->id); - if (expect_synthetic_inverse_frequencies) { - REQUIRE(scheduler.plan().transients.size() == 1); - REQUIRE(scheduler.plan().constant_initializations.size() == 1); - REQUIRE(dispatch.bindings[8].value == scheduler.plan().transients[0].value); - REQUIRE(scheduler.plan().constant_initializations[0].value == dispatch.bindings[8].value); - REQUIRE(scheduler.plan().constant_initializations[0].data.size() == - static_cast(kQwenFlashHeadSize / 2) * sizeof(float)); - } else { - REQUIRE(scheduler.plan().transients.empty()); - REQUIRE(scheduler.plan().constant_initializations.empty()); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + REQUIRE(output != nullptr); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 256); } - 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.initialization_commands.empty()); - REQUIRE(commands.commands.size() == 4); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.back().bindings.size() == 12); - REQUIRE(commands.commands.back().bindings[0].name == "positions"); - REQUIRE(commands.commands.back().bindings[1].name == "key_cache_indices"); - REQUIRE(commands.commands.back().bindings[2].name == "value_cache_indices"); - REQUIRE(commands.commands.back().bindings[3].name == "query_input"); - REQUIRE(commands.commands.back().bindings[4].name == "key_input"); - REQUIRE(commands.commands.back().bindings[5].name == "value_input"); - REQUIRE(commands.commands.back().bindings[6].name == "query_norm_weight"); - REQUIRE(commands.commands.back().bindings[7].name == "key_norm_weight"); - REQUIRE(commands.commands.back().bindings[8].name == "inverse_frequencies"); - REQUIRE(commands.commands.back().bindings[9].name == "query_output"); - REQUIRE(commands.commands.back().bindings[10].name == "key_cache"); - REQUIRE(commands.commands.back().bindings[11].name == "value_cache"); - REQUIRE(commands.constant_initializations.size() == (expect_synthetic_inverse_frequencies ? 1 : 0)); -} + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(bias != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * output = ggml_add(ctx, projection, bias); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); + } -static void run_qwen_attention_postprocess_dispatch_checks() { - 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 * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(bias != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * output = ggml_add(ctx, projection, bias); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_BF16, 33, 2048, + 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); + } { - constexpr int64_t token_count = 4; - constexpr int64_t query_head_count = 4; - constexpr int64_t key_value_head_count = 2; - constexpr int64_t cache_row_count = 16; - const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( - ctx, token_count, query_head_count, key_value_head_count, cache_row_count); - schedule_qwen_attention_postprocess_command(ctx, tensors, token_count, query_head_count, key_value_head_count, - cache_row_count); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 33); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(residual != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * output = ggml_add(ctx, projection, residual); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); } { - const QwenAttentionPostprocessTensors tensors = - build_qwen_attention_postprocess_graph(ctx, 4, 4, 2, 16, 0.00001f); - ggml::hrx::GraphImportResult imported = import_qwen_attention_postprocess_graph(ctx, tensors); - REQUIRE(!ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target())); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 33); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(bias != nullptr); + REQUIRE(residual != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * biased = ggml_add(ctx, projection, bias); + REQUIRE(biased != nullptr); + ggml_tensor * output = ggml_add(ctx, biased, residual); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 3, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD }, + { "input", "weight", "bias", "residual_input", "residual_output" }, false); } { - const QwenAttentionPostprocessTensors tensors = - build_qwen_attention_postprocess_graph(ctx, 4, 4, 2, 16, 0.000001f, false); - schedule_qwen_attention_postprocess_command(ctx, tensors, 4, 4, 2, 16, true); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 33); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(residual != nullptr); + REQUIRE(norm_weight != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * added = ggml_add(ctx, projection, residual); + REQUIRE(added != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, added, 0.000001f); + REQUIRE(rms != nullptr); + ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 4, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + { "input", "weight", "bias", "residual_input", "residual_output" }, true); } { - constexpr int64_t token_count = 13; - constexpr int64_t query_head_count = 32; - constexpr int64_t key_value_head_count = 4; - constexpr int64_t cache_row_count = 512; - const QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( - ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); - schedule_qwen_attention_postprocess_command(ctx, tensors, token_count, query_head_count, key_value_head_count, - cache_row_count, true); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 33); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 33); + ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(bias != nullptr); + REQUIRE(residual != nullptr); + REQUIRE(norm_weight != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + REQUIRE(projection != nullptr); + ggml_tensor * biased = ggml_add(ctx, projection, bias); + REQUIRE(biased != nullptr); + ggml_tensor * added = ggml_add(ctx, biased, residual); + REQUIRE(added != nullptr); + ggml_tensor * rms = ggml_rms_norm(ctx, added, 0.000001f); + REQUIRE(rms != nullptr); + ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); + REQUIRE(output != nullptr); + schedule_fused_matmul_postops_command( + ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, + 256, 5, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + { "input", "weight", "bias", "residual_input", "residual_output" }, true); } { - constexpr int64_t token_count = 13; - constexpr int64_t query_head_count = 32; - constexpr int64_t key_value_head_count = 4; - constexpr int64_t cache_row_count = 512; + struct FormatCase { + ggml_type type; + int64_t output_size; + }; + + const FormatCase cases[] = { + { GGML_TYPE_Q3_K, 128 }, + { GGML_TYPE_Q4_K, 128 }, + { GGML_TYPE_Q6_K, 128 }, + { GGML_TYPE_IQ3_S, 128 }, + { GGML_TYPE_IQ4_NL, 128 }, + { GGML_TYPE_IQ4_XS, 128 }, + { GGML_TYPE_Q8_0, 128 }, + { GGML_TYPE_Q8_1, 128 }, + { GGML_TYPE_F16, 128 }, + { GGML_TYPE_BF16, 128 }, + { GGML_TYPE_F32, 256 }, + }; + + for (const FormatCase & c : cases) { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, c.type, 2048, c.output_size); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, c.type, 2048, c.output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, c.type, c.type, 4, 2048, c.output_size); + } + } - const QwenAttentionPostprocessTensors first = build_qwen_attention_postprocess_graph( - ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); - const QwenAttentionPostprocessTensors second = build_qwen_attention_postprocess_graph( - ctx, token_count, query_head_count, key_value_head_count, cache_row_count, 0.000001f, false); + { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 2048); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 2048); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_GEGLU); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, 2, 640, 2048, + ggml::hrx::BinaryKind::GeGLU); + } - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, first.query_output); - ggml_build_forward_expand(graph, first.key_output); - ggml_build_forward_expand(graph, first.value_output); - ggml_build_forward_expand(graph, second.query_output); - ggml_build_forward_expand(graph, second.key_output); - ggml_build_forward_expand(graph, second.value_output); + { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 2048); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 2048); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(output != nullptr); + require_matmul_swiglu_falls_back_to_compilable_plan(ctx, output); + } - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); + { + ggml_tensor * packed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 3); + REQUIRE(packed != nullptr); + ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, false); + REQUIRE(output != nullptr); + schedule_packed_glu_command(ctx, output, ggml::hrx::BinaryKind::SwiGLU, 128, 3, false); + } - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 8); + { + ggml_tensor * packed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 3); + REQUIRE(packed != nullptr); + ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_GEGLU, true); + REQUIRE(output != nullptr); + schedule_packed_glu_command(ctx, output, ggml::hrx::BinaryKind::GeGLU, 128, 3, true); + } - 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() == 8); - REQUIRE(command_program_verifies(commands)); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 256, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * packed = ggml_mul_mat(ctx, weight, input); + REQUIRE(packed != nullptr); + ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, false); + REQUIRE(output != nullptr); + schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, ggml::hrx::BinaryKind::SwiGLU, + false); + } + + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 256, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * packed = ggml_mul_mat(ctx, weight, input); + REQUIRE(packed != nullptr); + ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, true); + REQUIRE(output != nullptr); + schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, ggml::hrx::BinaryKind::SwiGLU, + true); } { - constexpr int64_t token_count = 4; - constexpr int64_t query_head_count = 4; - constexpr int64_t key_value_head_count = 2; - constexpr int64_t cache_row_count = 16; - QwenAttentionPostprocessTensors tensors = build_qwen_attention_postprocess_graph( - ctx, token_count, query_head_count, key_value_head_count, cache_row_count); - ggml_tensor * flash_output = append_qwen_flash_attention_consumer(ctx, tensors, token_count, query_head_count, - key_value_head_count, cache_row_count); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, tensors.key_output); - ggml_build_forward_expand(graph, tensors.value_output); - ggml_build_forward_expand(graph, flash_output); + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, GGML_TYPE_IQ4_NL, GGML_TYPE_Q5_0, 2, 640, 256); + } - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); + { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, GGML_TYPE_Q4_K, GGML_TYPE_F16, 4, 2048, 128); + } - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::CommandPlan plan; - ggml::hrx::DispatchMatch match; - REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, - producer_index_for_tensor(imported.graph, tensors.query_reshape), match)); - append_match_to_plan(plan, match, covered_nodes); - REQUIRE(plan.valid()); - REQUIRE(plan.initialization_dispatches.size() == 2); + { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, GGML_TYPE_F32, GGML_TYPE_Q8_1, 4, 2048, 128); + } - const ggml::hrx::Dispatch & context_capture = plan.initialization_dispatches[0]; - const ggml::hrx::Dispatch & metadata = plan.initialization_dispatches[1]; - REQUIRE(kernel_name_for_id(context_capture.kernel.kernel_id) == - "qwen3_moe:qwen_attention_context_base_capture"); - REQUIRE(kernel_name_for_id(metadata.kernel.kernel_id) == "qwen3_moe:qwen_attention_metadata"); - REQUIRE(context_capture.bindings.size() == 2); - REQUIRE(metadata.bindings.size() == 5); - REQUIRE(context_capture.bindings[1].value == metadata.bindings[0].value); - REQUIRE(metadata.kernel.integer_parameters.at("token_count") == token_count); - REQUIRE(metadata.kernel.integer_parameters.at("context_capacity") == cache_row_count); + { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * gate_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + ggml_tensor * up_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(gate_input != nullptr); + REQUIRE(up_input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, gate_input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, up_input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); + REQUIRE(output != nullptr); + require_matmul_swiglu_falls_back(ctx, output); + } - const ggml::hrx::Value * positions_value = imported.graph.values().find_tensor(tensors.positions); - const ggml::hrx::Value * key_indices_value = imported.graph.values().find_tensor(tensors.key_cache_indices); - const ggml::hrx::Value * value_indices_value = imported.graph.values().find_tensor(tensors.value_cache_indices); - const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(tensors.mask); - REQUIRE(positions_value != nullptr); - REQUIRE(key_indices_value != nullptr); - REQUIRE(value_indices_value != nullptr); - REQUIRE(mask_value != nullptr); - REQUIRE(context_capture.bindings[0].value == positions_value->id); - REQUIRE(metadata.bindings[1].value == positions_value->id); - REQUIRE(metadata.bindings[2].value == key_indices_value->id); - REQUIRE(metadata.bindings[3].value == value_indices_value->id); - REQUIRE(metadata.bindings[4].value == mask_value->id); + for (const auto & variant : { + std::pair{ GGML_GLU_OP_GEGLU, ggml::hrx::BinaryKind::GeGLU }, + std::pair{ GGML_GLU_OP_REGLU, ggml::hrx::BinaryKind::RegLU }, + std::pair{ GGML_GLU_OP_GEGLU_ERF, ggml::hrx::BinaryKind::GeGLUErf }, + std::pair{ GGML_GLU_OP_GEGLU_QUICK, + ggml::hrx::BinaryKind::GeGLUQuick }, + }) { + ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(gate_weight != nullptr); + REQUIRE(up_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); + ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); + REQUIRE(gate != nullptr); + REQUIRE(up != nullptr); + ggml_tensor * output = ggml_glu_split(ctx, gate, up, variant.first); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 4, 2048, 128, variant.second, + GGML_OP_GLU); + } - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.initialization_commands.size() == 2); - REQUIRE(commands.initialization_commands[0].bindings[0].name == "positions"); - REQUIRE(commands.initialization_commands[0].bindings[1].name == "control"); - REQUIRE(commands.initialization_commands[1].bindings[0].name == "control"); - REQUIRE(commands.initialization_commands[1].bindings[4].name == "attention_mask"); - REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, context_capture.bindings[1].value) != - nullptr); + { + ggml_tensor * lhs_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * rhs_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(lhs_weight != nullptr); + REQUIRE(rhs_weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * lhs = ggml_mul_mat(ctx, lhs_weight, input); + ggml_tensor * rhs = ggml_mul_mat(ctx, rhs_weight, input); + REQUIRE(lhs != nullptr); + REQUIRE(rhs != nullptr); + ggml_tensor * output = ggml_sub(ctx, lhs, rhs); + REQUIRE(output != nullptr); + schedule_fused_matmul_swiglu_command(ctx, output, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 4, 2048, 128, + ggml::hrx::BinaryKind::Sub, GGML_OP_SUB); } - ggml_free(ctx); -} + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * output = ggml_silu(ctx, matmul); + REQUIRE(output != nullptr); + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + ggml::hrx::UnaryKind::Silu, 4, 2048, 128); + } -static void run_qwen_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); + { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * output = ggml_relu(ctx, matmul); + REQUIRE(output != nullptr); + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + ggml::hrx::UnaryKind::Relu, 4, 2048, 128); + } { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); REQUIRE(weight != nullptr); REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * output = ggml_gelu(ctx, matmul); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_dense_linear_q4k_f16_wmma", 4, 2048, 128); + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + ggml::hrx::UnaryKind::Gelu, 4, 2048, 128); } { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); REQUIRE(weight != nullptr); REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * output = ggml_gelu(ctx, matmul); REQUIRE(output != nullptr); - REQUIRE(!matmul_graph_is_supported(ctx, output)); + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + ggml::hrx::UnaryKind::Gelu, 2, 640, 128); } { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 2); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); REQUIRE(weight != nullptr); REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * relu = ggml_relu(ctx, matmul); + ggml_tensor * sqr = ggml_sqr(ctx, matmul); + REQUIRE(relu != nullptr); + REQUIRE(sqr != nullptr); + ggml_tensor * output = ggml_add(ctx, relu, sqr); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_dense_linear_q6k_f16_wmma", 2, 2048, 128); + require_matmul_root_matches_identity_single(ctx, output); } { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); REQUIRE(weight != nullptr); REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); + REQUIRE(matmul != nullptr); + ggml_tensor * output = ggml_softplus(ctx, matmul); REQUIRE(output != nullptr); - REQUIRE(!matmul_graph_is_supported(ctx, output)); + require_matmul_root_matches_identity_single(ctx, output); + } + + { + 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; + const std::string original_env_value = had_original_env ? original_env : ""; + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + + ggml_tensor * q4_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * q4_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(q4_weight != nullptr); + REQUIRE(q4_input != nullptr); + ggml_tensor * q4_output = ggml_mul_mat(ctx, q4_weight, q4_input); + REQUIRE(q4_output != nullptr); + schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + + ggml_tensor * q4_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); + ggml_tensor * q4_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(q4_decode_weight != nullptr); + REQUIRE(q4_decode_input != nullptr); + ggml_tensor * q4_decode_output = ggml_mul_mat(ctx, q4_decode_weight, q4_decode_input); + REQUIRE(q4_decode_output != nullptr); + schedule_single_matmul_command(ctx, q4_decode_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 2048, 128); + + for (ggml_type legacy_type : + { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1 }) { + ggml_tensor * legacy_weight = ggml_new_tensor_2d(ctx, legacy_type, 2048, 128); + ggml_tensor * legacy_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(legacy_weight != nullptr); + REQUIRE(legacy_input != nullptr); + ggml_tensor * legacy_output = ggml_mul_mat(ctx, legacy_weight, legacy_input); + REQUIRE(legacy_output != nullptr); + schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); + + ggml_tensor * legacy_decode_weight = ggml_new_tensor_2d(ctx, legacy_type, 2048, 128); + ggml_tensor * legacy_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(legacy_decode_weight != nullptr); + REQUIRE(legacy_decode_input != nullptr); + ggml_tensor * legacy_decode_output = ggml_mul_mat(ctx, legacy_decode_weight, legacy_decode_input); + REQUIRE(legacy_decode_output != nullptr); + schedule_single_matmul_command(ctx, legacy_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 2048, 128); + } + + ggml_tensor * q3_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q3_K, 2048, 128); + ggml_tensor * q3_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(q3_weight != nullptr); + REQUIRE(q3_input != nullptr); + ggml_tensor * q3_output = ggml_mul_mat(ctx, q3_weight, q3_input); + REQUIRE(q3_output != nullptr); + schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + + ggml_tensor * q3_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q3_K, 2048, 128); + ggml_tensor * q3_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(q3_decode_weight != nullptr); + REQUIRE(q3_decode_input != nullptr); + ggml_tensor * q3_decode_output = ggml_mul_mat(ctx, q3_decode_weight, q3_decode_input); + REQUIRE(q3_decode_output != nullptr); + schedule_single_matmul_command(ctx, q3_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, + 128); + + ggml_tensor * iq2_s_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ2_S, 2048, 640); + ggml_tensor * iq2_s_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 2); + REQUIRE(iq2_s_weight != nullptr); + 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_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); + REQUIRE(iq2_s_decode_weight != nullptr); + 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_vector_iq2_s_f32_f32", 1, + 2048, 640); + + ggml_tensor * iq3_s_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); + ggml_tensor * iq3_s_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(iq3_s_weight != nullptr); + 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_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); + REQUIRE(iq3_s_decode_weight != nullptr); + 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_vector_iq3_s_f32_f32", 1, + 2048, 128); + + ggml_tensor * iq4_nl_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); + ggml_tensor * iq4_nl_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(iq4_nl_weight != nullptr); + REQUIRE(iq4_nl_input != nullptr); + ggml_tensor * iq4_nl_output = ggml_mul_mat(ctx, iq4_nl_weight, iq4_nl_input); + REQUIRE(iq4_nl_output != nullptr); + schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); + + ggml_tensor * iq4_nl_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); + ggml_tensor * iq4_nl_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(iq4_nl_decode_weight != nullptr); + REQUIRE(iq4_nl_decode_input != nullptr); + ggml_tensor * iq4_nl_decode_output = ggml_mul_mat(ctx, iq4_nl_decode_weight, iq4_nl_decode_input); + REQUIRE(iq4_nl_decode_output != nullptr); + schedule_single_matmul_command(ctx, iq4_nl_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 2048, 128); + + ggml_tensor * iq4_nl_small_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 1024); + ggml_tensor * iq4_nl_small_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(iq4_nl_small_weight != nullptr); + REQUIRE(iq4_nl_small_input != nullptr); + ggml_tensor * iq4_nl_small_output = ggml_mul_mat(ctx, iq4_nl_small_weight, iq4_nl_small_input); + REQUIRE(iq4_nl_small_output != nullptr); + schedule_single_matmul_command(ctx, iq4_nl_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", + 2, 640, 1024); + + ggml_tensor * iq4_nl_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 1024); + ggml_tensor * iq4_nl_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); + REQUIRE(iq4_nl_small_decode_weight != nullptr); + REQUIRE(iq4_nl_small_decode_input != nullptr); + ggml_tensor * iq4_nl_small_decode_output = + ggml_mul_mat(ctx, iq4_nl_small_decode_weight, iq4_nl_small_decode_input); + REQUIRE(iq4_nl_small_decode_output != nullptr); + schedule_single_matmul_command(ctx, iq4_nl_small_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", + 1, 640, 1024); + + ggml_tensor * q5_0_small_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 256); + ggml_tensor * q5_0_small_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); + REQUIRE(q5_0_small_weight != nullptr); + REQUIRE(q5_0_small_input != nullptr); + ggml_tensor * q5_0_small_output = ggml_mul_mat(ctx, q5_0_small_weight, q5_0_small_input); + REQUIRE(q5_0_small_output != nullptr); + schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, + 640, 256); + + ggml_tensor * q5_0_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 256); + ggml_tensor * q5_0_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); + REQUIRE(q5_0_small_decode_weight != nullptr); + REQUIRE(q5_0_small_decode_input != nullptr); + ggml_tensor * q5_0_small_decode_output = ggml_mul_mat(ctx, q5_0_small_decode_weight, q5_0_small_decode_input); + REQUIRE(q5_0_small_decode_output != nullptr); + schedule_single_matmul_command(ctx, q5_0_small_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 640, 256); + + ggml_tensor * q6_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); + ggml_tensor * q6_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 2); + REQUIRE(q6_weight != nullptr); + REQUIRE(q6_input != nullptr); + ggml_tensor * q6_output = ggml_mul_mat(ctx, q6_weight, q6_input); + REQUIRE(q6_output != nullptr); + schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, + 128); + + ggml_tensor * q6_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); + ggml_tensor * q6_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(q6_decode_weight != nullptr); + REQUIRE(q6_decode_input != nullptr); + ggml_tensor * q6_decode_output = ggml_mul_mat(ctx, q6_decode_weight, q6_decode_input); + REQUIRE(q6_decode_output != nullptr); + schedule_single_matmul_command(ctx, q6_decode_output, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, + 2048, 128); + + ggml_tensor * q6_gemma_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 3840, 262208); + ggml_tensor * q6_gemma_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 1); + REQUIRE(q6_gemma_decode_weight != nullptr); + REQUIRE(q6_gemma_decode_input != nullptr); + ggml_tensor * q6_gemma_decode_output = ggml_mul_mat(ctx, q6_gemma_decode_weight, q6_gemma_decode_input); + REQUIRE(q6_gemma_decode_output != nullptr); + schedule_single_matmul_command(ctx, q6_gemma_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 3840, 262208); + + ggml_tensor * iq4_xs_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); + ggml_tensor * iq4_xs_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(iq4_xs_weight != nullptr); + REQUIRE(iq4_xs_input != nullptr); + ggml_tensor * iq4_xs_output = ggml_mul_mat(ctx, iq4_xs_weight, iq4_xs_input); + REQUIRE(iq4_xs_output != nullptr); + schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 2048, 128); + + ggml_tensor * iq4_xs_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); + ggml_tensor * iq4_xs_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(iq4_xs_decode_weight != nullptr); + REQUIRE(iq4_xs_decode_input != nullptr); + ggml_tensor * iq4_xs_decode_output = ggml_mul_mat(ctx, iq4_xs_decode_weight, iq4_xs_decode_input); + REQUIRE(iq4_xs_decode_output != nullptr); + schedule_single_matmul_command(ctx, iq4_xs_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 2048, 128); + + ggml_tensor * q8_0_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); + ggml_tensor * q8_0_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(q8_0_weight != nullptr); + REQUIRE(q8_0_input != nullptr); + ggml_tensor * q8_0_output = ggml_mul_mat(ctx, q8_0_weight, q8_0_input); + REQUIRE(q8_0_output != nullptr); + schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + + ggml_tensor * q8_0_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); + ggml_tensor * q8_0_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(q8_0_decode_weight != nullptr); + REQUIRE(q8_0_decode_input != nullptr); + ggml_tensor * q8_0_decode_output = ggml_mul_mat(ctx, q8_0_decode_weight, q8_0_decode_input); + REQUIRE(q8_0_decode_output != nullptr); + schedule_single_matmul_command(ctx, q8_0_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, + 128); + + ggml_tensor * q8_0_gemma_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 3840, 262208); + ggml_tensor * q8_0_gemma_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 1); + REQUIRE(q8_0_gemma_decode_weight != nullptr); + REQUIRE(q8_0_gemma_decode_input != nullptr); + ggml_tensor * q8_0_gemma_decode_output = ggml_mul_mat(ctx, q8_0_gemma_decode_weight, q8_0_gemma_decode_input); + REQUIRE(q8_0_gemma_decode_output != nullptr); + schedule_single_matmul_command(ctx, q8_0_gemma_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, + 3840, 262208); + + ggml_tensor * q8_0_unsafe_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 32768, 262144); + ggml_tensor * q8_0_unsafe_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 32768, 1); + REQUIRE(q8_0_unsafe_decode_weight != nullptr); + REQUIRE(q8_0_unsafe_decode_input != nullptr); + ggml_tensor * q8_0_unsafe_decode_output = + ggml_mul_mat(ctx, q8_0_unsafe_decode_weight, q8_0_unsafe_decode_input); + REQUIRE(q8_0_unsafe_decode_output != nullptr); + REQUIRE(!matmul_graph_is_supported(ctx, q8_0_unsafe_decode_output)); + + ggml_tensor * q8_1_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); + ggml_tensor * q8_1_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(q8_1_weight != nullptr); + REQUIRE(q8_1_input != nullptr); + ggml_tensor * q8_1_output = ggml_mul_mat(ctx, q8_1_weight, q8_1_input); + REQUIRE(q8_1_output != nullptr); + schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + + ggml_tensor * q8_1_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); + ggml_tensor * q8_1_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(q8_1_decode_weight != nullptr); + REQUIRE(q8_1_decode_input != nullptr); + ggml_tensor * q8_1_decode_output = ggml_mul_mat(ctx, q8_1_decode_weight, q8_1_decode_input); + REQUIRE(q8_1_decode_output != nullptr); + schedule_single_matmul_command(ctx, q8_1_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, + 128); + + ggml_tensor * f16_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); + ggml_tensor * f16_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(f16_weight != nullptr); + REQUIRE(f16_input != nullptr); + ggml_tensor * f16_output = ggml_mul_mat(ctx, f16_weight, f16_input); + REQUIRE(f16_output != nullptr); + schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + + ggml_tensor * f16_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); + ggml_tensor * f16_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(f16_decode_weight != nullptr); + REQUIRE(f16_decode_input != nullptr); + ggml_tensor * f16_decode_output = ggml_mul_mat(ctx, f16_decode_weight, f16_decode_input); + REQUIRE(f16_decode_output != nullptr); + schedule_single_matmul_command(ctx, f16_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, + 128); + + ggml_tensor * f32_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); + ggml_tensor * f32_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); + REQUIRE(f32_weight != nullptr); + REQUIRE(f32_input != nullptr); + ggml_tensor * f32_output = ggml_mul_mat(ctx, f32_weight, f32_input); + REQUIRE(f32_output != nullptr); + schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); + + ggml_tensor * f32_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); + ggml_tensor * f32_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); + REQUIRE(f32_decode_weight != nullptr); + REQUIRE(f32_decode_input != nullptr); + ggml_tensor * f32_decode_output = ggml_mul_mat(ctx, f32_decode_weight, f32_decode_input); + REQUIRE(f32_decode_output != nullptr); + schedule_single_matmul_command(ctx, f32_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, + 128); + + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); } { @@ -2298,7 +8696,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "qwen3_moe:qwen3_moe_dense_linear_q6k_f16_wmma", 1, 2048, 151936); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, 2048, + 151936); } { @@ -2330,7 +8729,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - REQUIRE(!matmul_graph_is_supported(ctx, output)); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, + 128); } { @@ -2354,8 +8754,8 @@ static void run_qwen_matmul_dispatch_checks() { } { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1024, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1024, 4); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 128); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); REQUIRE(weight != nullptr); REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); @@ -2366,6 +8766,413 @@ 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; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + for (ggml_type weight_type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + for (ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { + ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, 768, 256); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 768, 1); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, cache_type, 256, 17); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + ggml_tensor * output = ggml_set_rows(ctx, cache, ggml_mul_mat(ctx, weight, input), indices); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + auto imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == 1 || plan.dispatches.size() == 2); + const auto & dispatch = plan.dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + if (plan.dispatches.size() == 1) { + REQUIRE(kernel_name == "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32"); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", + std::to_string(matmul_weight_format_config(weight_type))); + } else { + REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + REQUIRE(kernel_name == "loom_libs:llm_attention_v_matmul_set_rows_legacy_f32_f32"); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", + weight_type == GGML_TYPE_Q4_K ? "44" : "46"); + require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); + REQUIRE(!dispatch.bindings[1].layout.empty()); + } + require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", + cache_type == GGML_TYPE_F16 ? "16" : "32"); + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); + } + } + ggml_free(ctx); +} + +static void run_prefill_value_projection_cache_dispatch_checks() { + static constexpr const char * kDisableFusionEnv = "GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"; + EnvironmentVariableGuard disable_fusion(kDisableFusionEnv); + disable_fusion.unset(); + + auto check_case = [&](ggml_type weight_type, int64_t input_size, int64_t output_size, int64_t token_count, + ggml_type cache_type, int64_t cache_row_count, bool publish_f16_alternate, + bool set_rows_available, bool disable, bool expect_fusion, bool observe_projection = false, + bool alias_weight = false) { + 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 * weight_storage = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); + REQUIRE(weight_storage != nullptr); + ggml_tensor * weight = + alias_weight ? ggml_view_2d(ctx, weight_storage, input_size, output_size, weight_storage->nb[1], 0) : + weight_storage; + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, cache_type, output_size, cache_row_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * value = ggml_reshape_3d(ctx, projection, 128, output_size / 128, token_count); + ggml_tensor * rows = ggml_reshape_2d(ctx, value, output_size, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(projection != nullptr); + REQUIRE(value != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + if (observe_projection) { + ggml_set_output(projection); + } + + 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::CommandPlan plan; + const ggml::hrx::Value * input_value = imported.graph.values().find_tensor(input); + REQUIRE(input_value != nullptr); + if (publish_f16_alternate) { + const size_t bytes = static_cast(input_size * token_count) * sizeof(ggml_fp16_t); + const ggml::hrx::ValueId alternate(static_cast(imported.graph.values().size())); + plan.transients.push_back({ alternate, "test.prefill_v_cache.f16", bytes, 256 }); + plan.constant_initializations.push_back( + { alternate, "test.prefill_v_cache.f16", 0, std::vector(bytes) }); + ggml::hrx::Status status; + REQUIRE(plan.metadata.append_alternate_value( + { input_value->id, alternate, GGML_TYPE_F16, bytes, "test.prefill_v_cache.f16" }, status)); + } + + std::vector covered_nodes(imported.graph.nodes().size(), false); + if (!set_rows_available) { + covered_nodes[producer_index_for_tensor(imported.graph, output)] = true; + } + if (disable) { + disable_fusion.set("1"); + } else { + disable_fusion.unset(); + } + + ggml::hrx::DispatchMatch match; + const size_t projection_index = producer_index_for_tensor(imported.graph, projection); + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, projection_index, match)); + REQUIRE(!match.dispatches.empty()); + const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); + const bool fused = + kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32"; + REQUIRE(fused == expect_fusion); + if (expect_fusion) { + REQUIRE(match.dispatches.size() == 1); + REQUIRE(match.covered_nodes.size() == 4); + REQUIRE(dispatch.bindings.size() == 4); + REQUIRE(dispatch.bindings[0].value == plan.transients.front().value); + REQUIRE(dispatch.bindings[0].length == plan.transients.front().size); + require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", "6"); + require_compile_parameter(dispatch, "llm.attention_qkv.input_size", "1536"); + require_compile_parameter(dispatch, "llm.attention_qkv.output_size", "256"); + require_compile_parameter(dispatch, "llm.attention_qkv.cache_row_count", "512"); + require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", "16"); + + append_match_to_plan(plan, match, covered_nodes, &imported.graph); + 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() == 1); + REQUIRE(commands.commands.front().bindings.size() == 4); + REQUIRE(commands.commands.front().bindings[0].name == "input"); + REQUIRE(commands.commands.front().bindings[1].name == "weight"); + REQUIRE(commands.commands.front().bindings[2].name == "indices"); + REQUIRE(commands.commands.front().bindings[3].name == "cache"); + } + ggml_free(ctx); + }; + + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, true); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, true, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, false, true, false, false); + check_case(GGML_TYPE_Q4_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 2048, 256, 64, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 512, 64, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 32, GGML_TYPE_F16, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F32, 512, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 1024, true, true, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, false, false, false); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false, true); + check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false, false, true); +} + +static void run_llama_attention_matmul_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 16 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t input_size = 2048; + constexpr int64_t head_size = 128; + constexpr int64_t head_count = 2; + constexpr int64_t output_size = head_size * head_count; + constexpr int64_t token_count = 7; + constexpr int64_t cache_row_count = 16; + + const ggml_type weight_types[] = { GGML_TYPE_F16, GGML_TYPE_BF16 }; + for (ggml_type weight_type : weight_types) { + { + 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); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(positions != nullptr); + REQUIRE(freqs != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * query = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); + ggml_tensor * output = ggml_rope_ext(ctx, query, positions, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(projection != nullptr); + REQUIRE(query != nullptr); + REQUIRE(output != nullptr); + schedule_fused_llama_attention_matmul_command( + ctx, output, "loom_libs:llm_attention_q_matmul_rope_tiled_f32_f32", weight_type, token_count, + input_size, output_size, head_size, head_count, 0, 0, 3, + { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE }, + { "input", "weight", "positions", "theta", "freq_factors", "output" }); + } + + { + 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); + ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, output_size, cache_row_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(positions != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * key = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); + ggml_tensor * rope = ggml_rope_ext(ctx, key, positions, freqs, head_size, GGML_ROPE_TYPE_NORMAL, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * rows = ggml_reshape_2d(ctx, rope, output_size, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(projection != nullptr); + REQUIRE(key != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + schedule_fused_llama_attention_matmul_command( + ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", weight_type, token_count, + input_size, output_size, head_size, head_count, cache_row_count, 16, 5, + { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, + { "input", "weight", "positions", "indices", "theta", "freq_factors", "cache" }); + } + + { + 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); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, output_size, cache_row_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + REQUIRE(cache != nullptr); + REQUIRE(indices != nullptr); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * value = ggml_reshape_3d(ctx, projection, head_size, head_count, token_count); + ggml_tensor * rows = ggml_reshape_2d(ctx, value, output_size, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(projection != nullptr); + REQUIRE(value != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(output != nullptr); + schedule_fused_llama_attention_matmul_command( + ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32", weight_type, token_count, + input_size, output_size, 0, 0, cache_row_count, 16, 4, + { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, + { "input", "weight", "indices", "cache" }); + } + } + + ggml_free(ctx); +} + static void schedule_qwen_router_top8_command(ggml_context * ctx, ggml_tensor * output, ggml_tensor * route_ids, @@ -2415,8 +9222,7 @@ static void schedule_qwen_router_top8_command(ggml_context * ctx, qwen_partition_table_size(token_count, expected_route_count, expected_expert_count); const bool uses_fused_prefill_expert_table_partition = token_count == 512 && expected_route_count == 8 && - route_ids->nb[1] / sizeof(int32_t) == static_cast(expected_route_count) && - expected_expert_count == 128; + route_ids->nb[1] / sizeof(int32_t) == static_cast(expected_route_count) && expected_expert_count == 128; REQUIRE(scheduler.plan().dispatches.size() == (uses_fused_prefill_expert_table_partition ? 2 : 3)); REQUIRE(scheduler.plan().transients.size() == 2); REQUIRE(scheduler.plan().constant_initializations.empty()); @@ -2473,7 +9279,7 @@ static void schedule_qwen_router_top8_command(ggml_context * ctx, REQUIRE(expert_table_partition_dispatch.bindings[3].length == sizeof(int32_t)); } else { const ggml::hrx::Dispatch & expert_table_dispatch = scheduler.plan().dispatches[1]; - REQUIRE(kernel_name_for_id(expert_table_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_build_expert_table"); + REQUIRE(kernel_name_for_id(expert_table_dispatch.kernel.kernel_id) == "loom_libs:ggml_moe_build_expert_table"); REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("token_count") == token_count); REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("route_count") == expected_route_count); REQUIRE(expert_table_dispatch.kernel.integer_parameters.at("route_stride") == @@ -2484,16 +9290,15 @@ static void schedule_qwen_router_top8_command(ggml_context * ctx, REQUIRE(expert_table_dispatch.bindings[0].length == route_id_length); REQUIRE(expert_table_dispatch.bindings[1].value == expert_table_transient.value); REQUIRE(expert_table_dispatch.bindings[1].length == expert_table_bytes); - require_compile_parameter(expert_table_dispatch, "qwen3_moe.routed_gate_up.expert_count", + require_compile_parameter(expert_table_dispatch, "ggml.moe_routing.expert_count", std::to_string(expected_expert_count)); - require_compile_parameter(expert_table_dispatch, "qwen3_moe.routed_gate_up.route_count", + require_compile_parameter(expert_table_dispatch, "ggml.moe_routing.route_count", std::to_string(expected_route_count)); - require_compile_parameter(expert_table_dispatch, "qwen3_moe.workload.token_capacity", - std::to_string(token_count)); + require_compile_parameter(expert_table_dispatch, "ggml.workload.token_capacity", std::to_string(token_count)); const ggml::hrx::Dispatch & partition_table_dispatch = scheduler.plan().dispatches[2]; REQUIRE(kernel_name_for_id(partition_table_dispatch.kernel.kernel_id) == - "qwen3_moe:qwen3_moe_build_expert_partition_table"); + "loom_libs:ggml_moe_build_expert_partition_table"); REQUIRE(partition_table_dispatch.kernel.integer_parameters.at("token_count") == token_count); REQUIRE(partition_table_dispatch.kernel.integer_parameters.at("route_count") == expected_route_count); REQUIRE(partition_table_dispatch.kernel.integer_parameters.at("expert_count") == expected_expert_count); @@ -2502,12 +9307,16 @@ static void schedule_qwen_router_top8_command(ggml_context * ctx, REQUIRE(partition_table_dispatch.bindings[0].length == expert_table_bytes); REQUIRE(partition_table_dispatch.bindings[1].value == partition_table_transient.value); REQUIRE(partition_table_dispatch.bindings[1].length == partition_table_bytes); - require_compile_parameter(partition_table_dispatch, "qwen3_moe.routed_gate_up.expert_count", + require_compile_parameter(partition_table_dispatch, "ggml.moe_routing.expert_count", std::to_string(expected_expert_count)); - require_compile_parameter(partition_table_dispatch, "qwen3_moe.routed_gate_up.route_count", + require_compile_parameter(partition_table_dispatch, "ggml.moe_routing.route_count", std::to_string(expected_route_count)); - require_compile_parameter(partition_table_dispatch, "qwen3_moe.workload.token_capacity", + require_compile_parameter(partition_table_dispatch, "ggml.workload.token_capacity", std::to_string(token_count)); + require_compile_parameter(partition_table_dispatch, "ggml.moe_routing.descriptor_expert_mask", "127"); + require_compile_parameter(partition_table_dispatch, "ggml.moe_routing.descriptor_partition_shift", "7"); + require_compile_parameter(partition_table_dispatch, "ggml.moe_routing.descriptor_row_count_shift", "13"); + require_compile_parameter(partition_table_dispatch, "ggml.moe_routing.partition_workgroup_size", "128"); } const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( @@ -2629,6 +9438,10 @@ static void schedule_manual_qwen_router_top8_command(ggml::hrx::Graph & graph, struct QwenRoutedGateUpTensors { ggml_tensor * route_ids = nullptr; ggml_tensor * route_weights = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * gate_weight = nullptr; + ggml_tensor * up_weight = nullptr; + ggml_tensor * down_weight = nullptr; ggml_tensor * gate = nullptr; ggml_tensor * up = nullptr; ggml_tensor * glu = nullptr; @@ -2655,27 +9468,27 @@ static QwenRoutedGateUpTensors build_qwen_routed_gate_up_graph(ggml_context * ct REQUIRE(tensors.route_weights != nullptr); REQUIRE(tensors.route_ids != nullptr); - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, 1, token_count); - ggml_tensor * gate_weight = + tensors.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, 1, token_count); + tensors.gate_weight = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, kQwenMoeHiddenSize, kQwenMoeIntermediateSize, kQwenRouterExpertCount); - ggml_tensor * up_weight = + tensors.up_weight = ggml_new_tensor_3d(ctx, up_weight_type, kQwenMoeHiddenSize, kQwenMoeIntermediateSize, kQwenRouterExpertCount); - REQUIRE(input != nullptr); - REQUIRE(gate_weight != nullptr); - REQUIRE(up_weight != nullptr); + REQUIRE(tensors.input != nullptr); + REQUIRE(tensors.gate_weight != nullptr); + REQUIRE(tensors.up_weight != nullptr); - tensors.gate = ggml_mul_mat_id(ctx, gate_weight, input, tensors.route_ids); - tensors.up = ggml_mul_mat_id(ctx, up_weight, input, tensors.route_ids); + tensors.gate = ggml_mul_mat_id(ctx, tensors.gate_weight, tensors.input, tensors.route_ids); + tensors.up = ggml_mul_mat_id(ctx, tensors.up_weight, tensors.input, tensors.route_ids); REQUIRE(tensors.gate != nullptr); REQUIRE(tensors.up != nullptr); tensors.glu = ggml_glu_split(ctx, tensors.gate, tensors.up, glu_op); REQUIRE(tensors.glu != nullptr); if (include_down) { - ggml_tensor * down_weight = ggml_new_tensor_3d(ctx, down_weight_type, kQwenMoeIntermediateSize, - kQwenMoeHiddenSize, kQwenRouterExpertCount); - REQUIRE(down_weight != nullptr); - tensors.output = ggml_mul_mat_id(ctx, down_weight, tensors.glu, tensors.route_ids); + tensors.down_weight = ggml_new_tensor_3d(ctx, down_weight_type, kQwenMoeIntermediateSize, kQwenMoeHiddenSize, + kQwenRouterExpertCount); + REQUIRE(tensors.down_weight != nullptr); + tensors.output = ggml_mul_mat_id(ctx, tensors.down_weight, tensors.glu, tensors.route_ids); REQUIRE(tensors.output != nullptr); } else { tensors.output = tensors.glu; @@ -2685,10 +9498,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) { @@ -2755,6 +9569,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); } @@ -2861,11 +9676,504 @@ static void append_qwen_routed_down_for_graph(ggml::hrx::Graph & gr REQUIRE(routed_down_alternate->byte_count == routed_down_transient.size); REQUIRE(dispatch.bindings[3].value == routed_down_transient.value); REQUIRE(dispatch.bindings[3].length == routed_down_transient.size); - require_compile_parameter(dispatch, "qwen3_moe.routed_down.input_size", "768"); - require_compile_parameter(dispatch, "qwen3_moe.routed_down.route_count", "8"); - require_compile_parameter(dispatch, "qwen3_moe.routed_down.expert_count", "128"); - require_compile_parameter(dispatch, "qwen3_moe.routed_down.output_size", "2048"); - require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(tensors.output->ne[2])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_f16_f16.input_size", "768"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_f16_f16.route_count", "8"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_f16_f16.expert_count", "128"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_f16_f16.output_size", "2048"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_f16_f16.weight_format", + tensors.down_weight->type == GGML_TYPE_Q4_K ? "4" : "6"); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(tensors.output->ne[2])); +} + +static std::string common_mul_mat_weight_format(ggml_type type) { + switch (type) { + case GGML_TYPE_Q3_K: + return "11"; + case GGML_TYPE_Q4_K: + return "4"; + case GGML_TYPE_Q5_K: + return "5"; + case GGML_TYPE_Q6_K: + return "6"; + case GGML_TYPE_IQ3_S: + return "21"; + case GGML_TYPE_IQ4_NL: + return "20"; + case GGML_TYPE_IQ4_XS: + return "23"; + case GGML_TYPE_Q8_0: + return "80"; + case GGML_TYPE_Q8_1: + return "81"; + case GGML_TYPE_F16: + return "16"; + case GGML_TYPE_BF16: + return "30"; + case GGML_TYPE_F32: + return "32"; + default: + return "0"; + } +} + +static const char * common_mul_mat_id_f32_kernel_name(int64_t token_count) { + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_input_f32_publish_f32"; +} + +static const char * common_mul_mat_id_postops_f32_kernel_name(int64_t token_count, bool has_rmsnorm) { + if (has_rmsnorm) { + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32"; + } + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_postops_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_input_f32_postops_publish_f32"; +} + +static const char * common_mul_mat_id_swiglu_f32_kernel_name(int64_t token_count) { + return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32" : + "loom_libs:ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32"; +} + +static ggml::hrx::DispatchMatch require_common_mul_mat_id_for_graph(ggml::hrx::Graph & graph, + ggml::hrx::CommandPlan & plan, + std::vector & covered_nodes, + ggml_tensor * output, + ggml_tensor * input, + ggml_tensor * weight, + ggml_tensor * route_ids) { + const ggml::hrx::Value * output_value = graph.values().find_tensor(output); + const ggml::hrx::Value * input_value = graph.values().find_tensor(input); + const ggml::hrx::Value * weight_value = graph.values().find_tensor(weight); + const ggml::hrx::Value * route_ids_value = graph.values().find_tensor(route_ids); + REQUIRE(output_value != nullptr); + REQUIRE(input_value != nullptr); + REQUIRE(weight_value != nullptr); + REQUIRE(route_ids_value != nullptr); + const ggml::hrx::CommandPlanMoeRoutingBundle * routing_bundle = + plan.metadata.find_moe_routing_bundle(route_ids_value->id); + REQUIRE(routing_bundle != nullptr); + + const size_t output_index = producer_index_for_tensor(graph, output); + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, output_index, match)); + REQUIRE(match.dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = match.dispatches[0]; + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == common_mul_mat_id_f32_kernel_name(output->ne[2])); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == output->ne[2]); + REQUIRE(dispatch.bindings.size() == 5); + REQUIRE(dispatch.bindings[0].value == input_value->id); + REQUIRE(dispatch.bindings[0].length == input_value->byte_count); + REQUIRE(dispatch.bindings[1].value == routing_bundle->expert_table); + REQUIRE(dispatch.bindings[1].length == qwen_expert_table_size(output->ne[2], weight->ne[2])); + REQUIRE(dispatch.bindings[2].value == routing_bundle->partition_table); + REQUIRE(dispatch.bindings[2].length == qwen_partition_table_size(output->ne[2], route_ids->ne[0], weight->ne[2])); + REQUIRE(dispatch.bindings[3].value == weight_value->id); + REQUIRE(dispatch.bindings[3].length == weight_value->byte_count); + REQUIRE(dispatch.bindings[4].value == output_value->id); + REQUIRE(dispatch.bindings[4].length == output_value->byte_count); + require_compile_parameter(dispatch, "ggml.mul_mat_id.input_size", std::to_string(weight->ne[0])); + require_compile_parameter(dispatch, "ggml.mul_mat_id.output_size", std::to_string(weight->ne[1])); + require_compile_parameter(dispatch, "ggml.mul_mat_id.expert_count", std::to_string(weight->ne[2])); + require_compile_parameter(dispatch, "ggml.mul_mat_id.route_count", std::to_string(route_ids->ne[0])); + require_compile_parameter(dispatch, "ggml.mul_mat_id.input_route_count", std::to_string(input->ne[1])); + require_compile_parameter(dispatch, "ggml.mul_mat_id.weight_format", common_mul_mat_weight_format(weight->type)); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(output->ne[2])); + return match; +} + +struct CommonMulMatIdSwiGLUTensors { + ggml_tensor * input = nullptr; + ggml_tensor * route_ids = nullptr; + ggml_tensor * gate_weight = nullptr; + ggml_tensor * up_weight = nullptr; + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + ggml_tensor * output = nullptr; +}; + +static CommonMulMatIdSwiGLUTensors build_common_mul_mat_id_swiglu_graph(ggml_context * ctx, + ggml_type gate_weight_type, + ggml_type up_weight_type, + ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, + int64_t input_size = 256, + int64_t output_size = 128, + int64_t token_count = 4) { + CommonMulMatIdSwiGLUTensors tensors; + constexpr int64_t route_count = 8; + constexpr int64_t input_route_count = 1; + constexpr int64_t expert_count = 16; + tensors.route_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, token_count); + tensors.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, input_size, input_route_count, token_count); + tensors.gate_weight = ggml_new_tensor_3d(ctx, gate_weight_type, input_size, output_size, expert_count); + tensors.up_weight = ggml_new_tensor_3d(ctx, up_weight_type, input_size, output_size, expert_count); + REQUIRE(tensors.route_ids != nullptr); + REQUIRE(tensors.input != nullptr); + REQUIRE(tensors.gate_weight != nullptr); + REQUIRE(tensors.up_weight != nullptr); + + tensors.gate = ggml_mul_mat_id(ctx, tensors.gate_weight, tensors.input, tensors.route_ids); + tensors.up = ggml_mul_mat_id(ctx, tensors.up_weight, tensors.input, tensors.route_ids); + REQUIRE(tensors.gate != nullptr); + REQUIRE(tensors.up != nullptr); + tensors.output = ggml_glu_split(ctx, tensors.gate, tensors.up, glu_op); + REQUIRE(tensors.output != nullptr); + return tensors; +} + +static ggml::hrx::GraphImportResult import_common_mul_mat_id_swiglu_graph(ggml_context * ctx, + const CommonMulMatIdSwiGLUTensors & tensors) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return imported; +} + +static void require_common_mul_mat_id_swiglu_match(ggml_context * ctx, + const CommonMulMatIdSwiGLUTensors & tensors, + ggml_tensor * root_projection) { + ggml::hrx::GraphImportResult imported = import_common_mul_mat_id_swiglu_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + const size_t root_index = producer_index_for_tensor(imported.graph, root_projection); + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, root_index, match)); + REQUIRE(match.dispatches.size() == 3); + REQUIRE(match.transients.size() == 2); + REQUIRE(match.metadata.generated_resources().size() == 2); + REQUIRE(match.metadata.moe_routing_bundles().size() == 1); + + const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.output->ne[2]); + REQUIRE(dispatch.bindings.size() == 6); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.input_size", + std::to_string(tensors.gate_weight->ne[0])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.output_size", + std::to_string(tensors.gate_weight->ne[1])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.expert_count", + std::to_string(tensors.gate_weight->ne[2])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.route_count", std::to_string(tensors.route_ids->ne[0])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.input_route_count", + std::to_string(tensors.input->ne[1])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.gate_weight_format", + common_mul_mat_weight_format(tensors.gate_weight->type)); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.up_weight_format", + common_mul_mat_weight_format(tensors.up_weight->type)); + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(tensors.output->ne[2])); + + const ggml::hrx::Value * input_value = imported.graph.values().find_tensor(tensors.input); + const ggml::hrx::Value * route_ids_value = imported.graph.values().find_tensor(tensors.route_ids); + const ggml::hrx::Value * gate_weight_value = imported.graph.values().find_tensor(tensors.gate_weight); + const ggml::hrx::Value * up_weight_value = imported.graph.values().find_tensor(tensors.up_weight); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(tensors.output); + REQUIRE(input_value != nullptr); + REQUIRE(route_ids_value != nullptr); + REQUIRE(gate_weight_value != nullptr); + REQUIRE(up_weight_value != nullptr); + REQUIRE(output_value != nullptr); + const ggml::hrx::CommandPlanMoeRoutingBundle * routing_bundle = + match.metadata.find_moe_routing_bundle(route_ids_value->id); + REQUIRE(routing_bundle != nullptr); + + REQUIRE(dispatch.bindings[0].value == input_value->id); + REQUIRE(dispatch.bindings[1].value == routing_bundle->expert_table); + REQUIRE(dispatch.bindings[2].value == routing_bundle->partition_table); + REQUIRE(dispatch.bindings[3].value == gate_weight_value->id); + REQUIRE(dispatch.bindings[4].value == up_weight_value->id); + REQUIRE(dispatch.bindings[5].value == output_value->id); + + append_match_to_plan(plan, match, covered_nodes, &imported.graph); + 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() == 3); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[2].bindings.size() == 6); + REQUIRE(commands.commands[2].bindings[0].name == "input"); + REQUIRE(commands.commands[2].bindings[1].name == "expert_table"); + REQUIRE(commands.commands[2].bindings[2].name == "partition_table"); + REQUIRE(commands.commands[2].bindings[3].name == "gate_weight"); + REQUIRE(commands.commands[2].bindings[4].name == "up_weight"); + REQUIRE(commands.commands[2].bindings[5].name == "output"); +} + +static void run_common_mul_mat_id_swiglu_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); + + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_Q4_K, GGML_TYPE_F16); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_F32, GGML_TYPE_Q8_1); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.up); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_BF16, GGML_TYPE_BF16); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_IQ4_XS, GGML_TYPE_IQ4_XS); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_Q3_K, GGML_TYPE_Q3_K); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_IQ3_S, GGML_TYPE_IQ3_S); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, GGML_GLU_OP_SWIGLU, 640, 128); + require_common_mul_mat_id_swiglu_match(ctx, tensors, tensors.gate); + } + { + const CommonMulMatIdSwiGLUTensors tensors = + build_common_mul_mat_id_swiglu_graph(ctx, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_GLU_OP_GEGLU); + ggml::hrx::GraphImportResult imported = import_common_mul_mat_id_swiglu_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, match)); + REQUIRE(kernel_name_for_id(match.dispatches.back().kernel.kernel_id) == + common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); + } + + ggml_free(ctx); +} + +struct CommonMulMatIdPostOpsTensors { + ggml_tensor * input = nullptr; + ggml_tensor * weight = nullptr; + ggml_tensor * route_ids = nullptr; + ggml_tensor * projection = nullptr; + ggml_tensor * bias = nullptr; + ggml_tensor * residual_input = nullptr; + ggml_tensor * residual_output = nullptr; + ggml_tensor * rms_output = nullptr; + ggml_tensor * norm_weight = nullptr; + ggml_tensor * normalized_output = nullptr; +}; + +static int64_t common_mul_mat_id_partition_descriptor_capacity(const CommonMulMatIdPostOpsTensors & tensors) { + const int64_t token_count = tensors.projection->ne[2]; + const int64_t route_count = tensors.route_ids->ne[0]; + const int64_t expert_count = tensors.weight->ne[2]; + const int64_t assignment_count = token_count * route_count; + const int64_t assignment_partition_count = (assignment_count + 31) / 32; + return assignment_partition_count + expert_count; +} + +static CommonMulMatIdPostOpsTensors build_common_mul_mat_id_postops_graph(ggml_context * ctx, + bool include_bias, + bool include_residual, + bool include_rmsnorm, + ggml_type weight_type = GGML_TYPE_Q4_K, + int64_t token_count = 4) { + CommonMulMatIdPostOpsTensors tensors; + constexpr int64_t route_count = 8; + constexpr int64_t input_size = 256; + constexpr int64_t output_size = 128; + constexpr int64_t expert_count = 16; + tensors.route_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, token_count); + tensors.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, input_size, 1, token_count); + tensors.weight = ggml_new_tensor_3d(ctx, weight_type, input_size, output_size, expert_count); + REQUIRE(tensors.route_ids != nullptr); + REQUIRE(tensors.input != nullptr); + REQUIRE(tensors.weight != nullptr); + + tensors.projection = ggml_mul_mat_id(ctx, tensors.weight, tensors.input, tensors.route_ids); + REQUIRE(tensors.projection != nullptr); + + ggml_tensor * current = tensors.projection; + if (include_bias) { + tensors.bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_size); + REQUIRE(tensors.bias != nullptr); + current = ggml_add(ctx, current, tensors.bias); + REQUIRE(current != nullptr); + } + if (include_residual) { + tensors.residual_input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, output_size, route_count, token_count); + REQUIRE(tensors.residual_input != nullptr); + current = ggml_add(ctx, current, tensors.residual_input); + REQUIRE(current != nullptr); + } + tensors.residual_output = current; + + if (include_rmsnorm) { + REQUIRE(include_residual); + tensors.norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_size); + REQUIRE(tensors.norm_weight != nullptr); + tensors.rms_output = ggml_rms_norm(ctx, tensors.residual_output, 0.000001f); + REQUIRE(tensors.rms_output != nullptr); + tensors.normalized_output = ggml_mul(ctx, tensors.rms_output, tensors.norm_weight); + REQUIRE(tensors.normalized_output != nullptr); + } else { + tensors.normalized_output = tensors.residual_output; + } + return tensors; +} + +static ggml::hrx::GraphImportResult import_common_mul_mat_id_postops_graph( + ggml_context * ctx, + const CommonMulMatIdPostOpsTensors & tensors) { + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, tensors.normalized_output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + return imported; +} + +static void require_common_mul_mat_id_postops_match(ggml_context * ctx, + const CommonMulMatIdPostOpsTensors & tensors, + const char * expected_kernel_name, + bool expect_bias, + bool expect_residual, + bool expect_rmsnorm) { + ggml::hrx::GraphImportResult imported = import_common_mul_mat_id_postops_graph(ctx, tensors); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + const size_t projection_index = producer_index_for_tensor(imported.graph, tensors.projection); + ggml::hrx::DispatchMatch match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, projection_index, match)); + REQUIRE(match.dispatches.size() == 3); + REQUIRE(match.transients.size() == 2); + + const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.projection->ne[2]); + REQUIRE(dispatch.bindings.size() == (expect_rmsnorm ? 10 : 7)); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.input_size", std::to_string(tensors.weight->ne[0])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.output_size", std::to_string(tensors.weight->ne[1])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.expert_count", std::to_string(tensors.weight->ne[2])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.route_count", + std::to_string(tensors.route_ids->ne[0])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.input_route_count", + std::to_string(tensors.input->ne[1])); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.weight_format", + common_mul_mat_weight_format(tensors.weight->type)); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.has_bias", expect_bias ? "1" : "0"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.has_residual", expect_residual ? "1" : "0"); + + const ggml::hrx::Value * input_value = imported.graph.values().find_tensor(tensors.input); + const ggml::hrx::Value * weight_value = imported.graph.values().find_tensor(tensors.weight); + const ggml::hrx::Value * route_ids_value = imported.graph.values().find_tensor(tensors.route_ids); + const ggml::hrx::Value * residual_output_value = imported.graph.values().find_tensor(tensors.residual_output); + const ggml::hrx::Value * normalized_output_value = imported.graph.values().find_tensor(tensors.normalized_output); + REQUIRE(input_value != nullptr); + REQUIRE(weight_value != nullptr); + REQUIRE(route_ids_value != nullptr); + REQUIRE(residual_output_value != nullptr); + REQUIRE(normalized_output_value != nullptr); + const ggml::hrx::CommandPlanMoeRoutingBundle * routing_bundle = + match.metadata.find_moe_routing_bundle(route_ids_value->id); + REQUIRE(routing_bundle != nullptr); + + REQUIRE(dispatch.bindings[0].value == input_value->id); + REQUIRE(dispatch.bindings[1].value == routing_bundle->expert_table); + REQUIRE(dispatch.bindings[2].value == routing_bundle->partition_table); + REQUIRE(dispatch.bindings[3].value == weight_value->id); + if (expect_bias) { + const ggml::hrx::Value * bias_value = imported.graph.values().find_tensor(tensors.bias); + REQUIRE(bias_value != nullptr); + REQUIRE(dispatch.bindings[4].value == bias_value->id); + } + if (expect_residual) { + const ggml::hrx::Value * residual_input_value = imported.graph.values().find_tensor(tensors.residual_input); + REQUIRE(residual_input_value != nullptr); + REQUIRE(dispatch.bindings[5].value == residual_input_value->id); + } + REQUIRE(dispatch.bindings[6].value == residual_output_value->id); + + if (expect_rmsnorm) { + const ggml::hrx::Value * norm_weight_value = imported.graph.values().find_tensor(tensors.norm_weight); + REQUIRE(norm_weight_value != nullptr); + REQUIRE(dispatch.bindings[7].value == norm_weight_value->id); + REQUIRE(dispatch.bindings[8].value == normalized_output_value->id); + REQUIRE(match.completion_counter_requests.size() == 1); + REQUIRE(match.completion_counter_requests.front().count == + common_mul_mat_id_partition_descriptor_capacity(tensors)); + REQUIRE(dispatch.bindings[9].value == match.completion_counter_requests.front().value); + require_compile_parameter(dispatch, "ggml.mul_mat_id_postops.rms_epsilon", "9.99999997e-07"); + } else { + REQUIRE(match.completion_counter_requests.empty()); + } + + append_match_to_plan(plan, match, covered_nodes, &imported.graph); + 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() == 3); +} + +static void run_common_mul_mat_id_postops_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); + + { + const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, false, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, false, + false); + } + { + const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), false, true, + false); + } + { + const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, false); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, true, + false); + } + { + const CommonMulMatIdPostOpsTensors tensors = + build_common_mul_mat_id_postops_graph(ctx, true, true, false, GGML_TYPE_BF16); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, true, + false); + } + { + const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, true); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), false, true, + true); + } + { + const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, true); + require_common_mul_mat_id_postops_match( + ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), true, true, true); + } + + ggml_free(ctx); } static void append_qwen_weighted_reduce_for_graph(ggml::hrx::Graph & graph, @@ -2882,32 +10190,39 @@ static void append_qwen_weighted_reduce_for_graph(ggml::hrx::Graph & const ggml::hrx::Value * route_weights_value = graph.values().find_tensor(tensors.route_weights); const ggml::hrx::Value * routed_output_value = graph.values().find_tensor(tensors.output); + const ggml::hrx::Value * hidden_state_value = graph.values().find_tensor(tensors.hidden_state); const ggml::hrx::Value * residual_value = graph.values().find_tensor(tensors.residual); REQUIRE(route_weights_value != nullptr); REQUIRE(routed_output_value != nullptr); + REQUIRE(hidden_state_value != nullptr); REQUIRE(residual_value != nullptr); const ggml::hrx::Dispatch & dispatch = plan.dispatches.back(); REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.output->ne[2]); - REQUIRE(dispatch.bindings.size() == (tensors.next_output != nullptr ? 5 : 3)); + REQUIRE(dispatch.bindings.size() == (tensors.next_output != nullptr ? 5 : 4)); REQUIRE(dispatch.bindings[0].value == route_weights_value->id); REQUIRE(dispatch.bindings[0].length == route_weights_value->byte_count); REQUIRE(dispatch.bindings[1].value == plan.metadata.alternate_values().back().alternate_value); REQUIRE(dispatch.bindings[1].length == qwen_routed_down_f16_output_size(tensors.output->ne[2])); - REQUIRE(dispatch.bindings[2].value == residual_value->id); - REQUIRE(dispatch.bindings[2].length == residual_value->byte_count); require_compile_parameter(dispatch, "qwen3_moe.routed_down.route_count", "8"); require_compile_parameter(dispatch, "qwen3_moe.routed_down.output_size", "2048"); require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(tensors.output->ne[2])); if (tensors.next_output != nullptr) { + REQUIRE(dispatch.bindings[2].value == residual_value->id); + REQUIRE(dispatch.bindings[2].length == residual_value->byte_count); const ggml::hrx::Value * next_output_value = graph.values().find_tensor(tensors.next_output); REQUIRE(next_output_value != nullptr); REQUIRE(dispatch.bindings[4].value == next_output_value->id); REQUIRE(dispatch.bindings[4].length == next_output_value->byte_count); require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", "2048"); require_compile_parameter(dispatch, "qwen3_moe.model.rms_epsilon", "0.000001"); + } else { + REQUIRE(dispatch.bindings[2].value == hidden_state_value->id); + REQUIRE(dispatch.bindings[2].length == hidden_state_value->byte_count); + REQUIRE(dispatch.bindings[3].value == residual_value->id); + REQUIRE(dispatch.bindings[3].length == residual_value->byte_count); } } @@ -2918,11 +10233,78 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); + { + constexpr int64_t token_count = 4; + QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count, GGML_GLU_OP_GEGLU); + 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); + ggml::hrx::DispatchMatch gate_match = require_common_mul_mat_id_for_graph( + imported.graph, plan, covered_nodes, tensors.gate, tensors.input, tensors.gate_weight, tensors.route_ids); + append_match_to_plan(plan, gate_match, covered_nodes, &imported.graph); + + REQUIRE(plan.dispatches.size() == 4); + REQUIRE(plan.transients.size() == 2); + 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() == 4); + REQUIRE(command_program_verifies(commands)); + REQUIRE(commands.commands[3].bindings.size() == 5); + REQUIRE(commands.commands[3].bindings[0].name == "input"); + REQUIRE(commands.commands[3].bindings[1].name == "expert_table"); + REQUIRE(commands.commands[3].bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[3].bindings[2].name == "partition_table"); + REQUIRE(commands.commands[3].bindings[2].origin == ggml::hrx::CommandBindingOrigin::Transient); + REQUIRE(commands.commands[3].bindings[3].name == "weight"); + REQUIRE(commands.commands[3].bindings[4].name == "output"); + } + { constexpr int64_t token_count = 4; const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); - REQUIRE(imported.graph.nodes().size() == 14); + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); + 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() == 4); + REQUIRE(plan.transients.size() == 2); + 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() == 4); + REQUIRE(commands.commands[3].bindings[0].name == "input"); + REQUIRE(commands.commands[3].bindings[2].name == "partition_table"); + REQUIRE(commands.commands[3].bindings[3].name == "weight"); + REQUIRE(commands.commands[3].bindings[4].name == "output"); + } + + { + constexpr int64_t token_count = 4; + const QwenRoutedGateUpTensors tensors = + build_qwen_routed_gate_up_graph(ctx, token_count, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, true, GGML_TYPE_BF16); + 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); + 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() == 4); + 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() == 4); + } + + { + constexpr int64_t token_count = 4; + 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); 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); @@ -2931,7 +10313,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); @@ -2982,8 +10364,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(plan.metadata.alternate_values().front().byte_count == f16_output_transient.size); const ggml::hrx::Dispatch & dispatch = plan.dispatches.back(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma"); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_id_swiglu_f16_f16_wmma"); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); REQUIRE(dispatch.bindings.size() == 6); REQUIRE(dispatch.bindings[1].value == routing_bundle->expert_table); @@ -2992,11 +10373,12 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(dispatch.bindings[2].length == qwen_partition_table_size(token_count)); REQUIRE(dispatch.bindings[5].value == f16_output_transient.value); REQUIRE(dispatch.bindings[5].length == f16_output_transient.size); - require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.input_size", "2048"); - require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.expert_count", "128"); - require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.route_count", "8"); - require_compile_parameter(dispatch, "qwen3_moe.routed_gate_up.output_size", "768"); - require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(token_count)); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu_f16_f16.input_size", "2048"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu_f16_f16.expert_count", "128"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu_f16_f16.route_count", "8"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu_f16_f16.output_size", "768"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu_f16_f16.gate_weight_format", "4"); + require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu_f16_f16.up_weight_format", "4"); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -3021,7 +10403,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); @@ -3033,7 +10416,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(plan.dispatches.size() == 4); REQUIRE(plan.transients.size() == 3); REQUIRE(kernel_name_for_id(plan.dispatches.back().kernel.kernel_id) == - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma"); + "loom_libs:ggml_mul_mat_id_swiglu_f16_f16_wmma"); REQUIRE(plan.dispatches.back().kernel.integer_parameters.at("token_count") == 1); const ggml::hrx::CommandProgram commands = @@ -3051,12 +10434,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()); @@ -3126,17 +10510,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, - "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + 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); } { @@ -3148,7 +10535,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { 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, - "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); @@ -3159,12 +10546,14 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 6); REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.back().bindings.size() == 3); + REQUIRE(commands.commands.back().bindings.size() == 4); REQUIRE(commands.commands.back().bindings[0].name == "route_weights"); REQUIRE(commands.commands.back().bindings[1].name == "routed_output"); REQUIRE(commands.commands.back().bindings[1].origin == ggml::hrx::CommandBindingOrigin::Transient); - REQUIRE(commands.commands.back().bindings[2].name == "output"); - REQUIRE(commands.commands.back().bindings[2].access == ggml::hrx::ResourceAccess::ReadWrite); + REQUIRE(commands.commands.back().bindings[2].name == "residual_input"); + REQUIRE(commands.commands.back().bindings[2].access == ggml::hrx::ResourceAccess::Read); + REQUIRE(commands.commands.back().bindings[3].name == "output"); + REQUIRE(commands.commands.back().bindings[3].access == ggml::hrx::ResourceAccess::Write); } { @@ -3176,7 +10565,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { 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, - "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); @@ -3188,6 +10577,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); @@ -3197,7 +10646,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { 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, - "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32"); @@ -3240,7 +10689,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { 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, - "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped"); + "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; @@ -3251,7 +10700,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(kernel_name_for_id(plan.dispatches.back().kernel.kernel_id) == "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); REQUIRE(weighted_match.value_aliases.empty()); - REQUIRE(plan.dispatches.back().bindings.size() == 3); + REQUIRE(plan.dispatches.back().bindings.size() == 4); const ggml::hrx::Value * hidden_state_value = imported.graph.values().find_tensor(tensors.hidden_state); const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(tensors.residual); REQUIRE(hidden_state_value != nullptr); @@ -3262,14 +10711,15 @@ 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); 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, - "qwen3_moe:qwen3_moe_routed_down_q4k_f16_wmma_grouped"); + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); REQUIRE(plan.dispatches.size() == 5); REQUIRE(plan.transients.size() == 4); @@ -3287,7 +10737,20 @@ static void run_qwen_routed_gate_up_dispatch_checks() { std::vector covered_nodes(imported.graph.nodes().size(), false); const ggml::hrx::CommandPlan empty_plan; ggml::hrx::DispatchMatch gate_up_match; - REQUIRE(!match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(gate_up_match.dispatches.size() == 3); + REQUIRE(gate_up_match.transients.size() == 2); + REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == + "loom_libs:ggml_moe_build_expert_table"); + REQUIRE(kernel_name_for_id(gate_up_match.dispatches[1].kernel.kernel_id) == + "loom_libs:ggml_moe_build_expert_partition_table"); + REQUIRE(kernel_name_for_id(gate_up_match.dispatches[2].kernel.kernel_id) == + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + REQUIRE(gate_up_match.dispatches[2].bindings.size() == 6); + REQUIRE(gate_up_match.dispatches[2].bindings[1].value == gate_up_match.transients[0].value); + REQUIRE(gate_up_match.dispatches[2].bindings[2].value == gate_up_match.transients[1].value); + REQUIRE(gate_up_match.metadata.generated_resources().size() == 2); + REQUIRE(gate_up_match.metadata.moe_routing_bundles().size() == 1); } { @@ -3297,7 +10760,10 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); ggml::hrx::DispatchMatch gate_up_match; - REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(gate_up_match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == + common_mul_mat_id_f32_kernel_name(tensors.gate->ne[2])); } { @@ -3308,7 +10774,11 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); ggml::hrx::DispatchMatch gate_up_match; - REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(gate_up_match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + REQUIRE(gate_up_match.dispatches[0].bindings.size() == 6); } { @@ -3319,7 +10789,11 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); const size_t gate_index = producer_index_for_tensor(imported.graph, tensors.gate); ggml::hrx::DispatchMatch gate_up_match; - REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); + REQUIRE(gate_up_match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == + common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + REQUIRE(gate_up_match.dispatches[0].bindings.size() == 6); } { @@ -3329,7 +10803,10 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); const size_t down_index = producer_index_for_tensor(imported.graph, tensors.output); ggml::hrx::DispatchMatch down_match; - REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); + REQUIRE(down_match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(down_match.dispatches[0].kernel.kernel_id) == + common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); } { @@ -3339,12 +10816,98 @@ 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); - const size_t down_index = producer_index_for_tensor(imported.graph, tensors.output); - ggml::hrx::DispatchMatch down_match; - REQUIRE(!match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); - } + const size_t down_index = producer_index_for_tensor(imported.graph, tensors.output); + ggml::hrx::DispatchMatch down_match; + REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); + REQUIRE(down_match.dispatches.size() == 1); + REQUIRE(kernel_name_for_id(down_match.dispatches[0].kernel.kernel_id) == + common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); + } + + 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"); - ggml_free(ctx); + 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() { @@ -3519,16 +11082,6 @@ static void run_alias_value_import_checks() { ggml::hrx::DispatchScheduler internal_scheduler; REQUIRE(internal_scheduler.schedule_graph(internal_imported.graph, test_dispatch_target())); - REQUIRE(internal_scheduler.plan().valid()); - REQUIRE(internal_scheduler.plan().dispatches.size() == 1); - const ggml::hrx::CommandProgram internal_commands = ggml::hrx::build_command_program( - internal_imported.graph, internal_scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(internal_commands.valid()); - REQUIRE(internal_commands.commands.size() == 1); - REQUIRE(internal_commands.commands[0].bindings[0].value == internal_view_value->id); - REQUIRE(internal_commands.commands[0].bindings[0].offset == 0); - REQUIRE(internal_commands.commands[0].bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); - REQUIRE(ggml::hrx::find_transient_allocation(internal_commands.transients, internal_view_value->id) == nullptr); REQUIRE(internal_imported.graph.values().bind_buffer( internal_source_value->id, { dummy_hrx_buffer(0x5000), 256, internal_source_value->byte_count })); @@ -3544,11 +11097,6 @@ static void run_alias_value_import_checks() { REQUIRE(internal_view_binding->buffer == dummy_hrx_buffer(0x5000)); REQUIRE(internal_view_binding->offset == 256 + internal_source->nb[1]); REQUIRE(internal_view_binding->length == internal_view_value->byte_count); - const ggml::hrx::ResolvedCommandProgram internal_resolved = - ggml::hrx::resolve_command_program_bindings(internal_commands, internal_bindings); - REQUIRE(internal_resolved.valid()); - REQUIRE(internal_resolved.commands[0].bindings[0].ref.buffer == dummy_hrx_buffer(0x5000)); - REQUIRE(internal_resolved.commands[0].bindings[0].ref.offset == 256 + internal_source->nb[1]); ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 4); ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8, 2); @@ -3720,7 +11268,7 @@ static void run_layout_alias_scheduler_elision_checks() { const ggml::hrx::Value * reshaped_value = imported.graph.values().find_tensor(reshaped); REQUIRE(sum_value != nullptr); REQUIRE(reshaped_value != nullptr); - REQUIRE(sum_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(sum_value->kind == ggml::hrx::ValueKind::External); REQUIRE(reshaped_value->kind == ggml::hrx::ValueKind::External); REQUIRE(imported.graph.values().same_storage(sum_value->id, reshaped_value->id)); REQUIRE(reshaped_value->storage_root == sum_value->id); @@ -3738,15 +11286,280 @@ static void run_layout_alias_scheduler_elision_checks() { REQUIRE(commands.commands.size() == 1); REQUIRE(commands.commands[0].bindings.size() == 3); REQUIRE(commands.commands[0].bindings[2].value == sum_value->id); - REQUIRE(commands.commands[0].bindings[2].origin == ggml::hrx::CommandBindingOrigin::Transient); - REQUIRE(commands.transients.allocations.size() == 1); - REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, sum_value->id) != nullptr); + REQUIRE(commands.commands[0].bindings[2].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.transients.allocations.empty()); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, sum_value->id) == nullptr); REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, reshaped_value->id) == nullptr); REQUIRE(command_program_verifies(commands)); ggml_free(ctx); } +static void run_zero_output_scheduler_elision_checks() { + ggml_init_params params = {}; + params.mem_size = 512 * 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, 3840, 1); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 0); + REQUIRE(input != nullptr); + REQUIRE(ids != nullptr); + ggml_tensor * rows = ggml_get_rows(ctx, input, ids); + REQUIRE(rows != nullptr); + REQUIRE(rows->ne[0] == 3840); + REQUIRE(rows->ne[1] == 0); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, rows); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + const ggml::hrx::Value * rows_value = imported.graph.values().find_tensor(rows); + REQUIRE(rows_value != nullptr); + REQUIRE(rows_value->element_count == 0); + REQUIRE(rows_value->byte_count == 0); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.empty()); + + 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.empty()); + REQUIRE(command_program_verifies(commands)); + + ggml::hrx::GraphProgramCache cache; + ggml::hrx::GraphProgramLookup lookup = + cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.match.external_bindings.size() == 1); + REQUIRE(lookup.match.external_bindings[0].tensor == input); + + const auto * ids_value = lookup.program->graph().values().find_tensor(ids); + REQUIRE(ids_value != nullptr); + ggml::hrx::Command referenced; + referenced.bindings.push_back({ "ids", ids_value->id }); + for (auto * list : + { &lookup.program->commands().initialization_commands, &lookup.program->commands().commands }) { + list->push_back(referenced); + const auto match = lookup.program->match_current_graph(*graph); + REQUIRE(match.valid()); + REQUIRE(match.external_bindings.size() == 2); + REQUIRE(std::any_of(match.external_bindings.begin(), match.external_bindings.end(), + [&](const auto & binding) { return binding.tensor == ids; })); + list->clear(); + const auto unreferenced = lookup.program->match_current_graph(*graph); + REQUIRE(unreferenced.valid()); + REQUIRE(unreferenced.external_bindings.size() == 1); + for (const auto origin : + { ggml::hrx::CommandBindingOrigin::Transient, ggml::hrx::CommandBindingOrigin::ProgramConstant }) { + ggml::hrx::Command unrelated; + unrelated.bindings.push_back({ "negative", ggml::hrx::ValueId(-1) }); + unrelated.bindings.push_back({ "large", ggml::hrx::ValueId(INT32_MAX) }); + unrelated.bindings.push_back({ "ids", ids_value->id, origin }); + list->push_back(unrelated); + const auto ignored = lookup.program->match_current_graph(*graph); + REQUIRE(ignored.valid()); + REQUIRE(ignored.external_bindings.size() == 1); + REQUIRE(ignored.external_bindings[0].tensor == input); + list->clear(); + } + } + } + + { + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 1); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 0); + REQUIRE(input != nullptr); + REQUIRE(ids != nullptr); + ggml_tensor * rows = ggml_get_rows(ctx, input, ids); + ggml_tensor * norm = ggml_rms_norm(ctx, rows, 0.000001f); + REQUIRE(rows != nullptr); + REQUIRE(norm != nullptr); + REQUIRE(norm->ne[0] == 3840); + REQUIRE(norm->ne[1] == 0); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, norm); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_RMS_NORM); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.empty()); + + ggml::hrx::GraphProgramCache cache; + ggml::hrx::GraphProgramLookup lookup = + cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.match.external_bindings.size() == 1); + REQUIRE(lookup.match.external_bindings[0].tensor == input); + } + + { + ggml_tensor * lhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 0); + ggml_tensor * rhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 0); + REQUIRE(lhs != nullptr); + REQUIRE(rhs != nullptr); + ggml_tensor * sum = ggml_add(ctx, lhs, rhs); + REQUIRE(sum != nullptr); + REQUIRE(sum->ne[0] == 3840); + REQUIRE(sum->ne[1] == 0); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, sum); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.empty()); + + ggml::hrx::GraphProgramCache cache; + ggml::hrx::GraphProgramLookup lookup = + cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.match.external_bindings.empty()); + } + + { + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 0); + REQUIRE(input != nullptr); + ggml_tensor * scaled = ggml_scale(ctx, input, 2.0f); + REQUIRE(scaled != nullptr); + REQUIRE(scaled->ne[0] == 3840); + REQUIRE(scaled->ne[1] == 0); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, scaled); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_SCALE); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.empty()); + + ggml::hrx::GraphProgramCache cache; + ggml::hrx::GraphProgramLookup lookup = + cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.match.external_bindings.empty()); + } + + { + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 1); + REQUIRE(input != nullptr); + ggml_tensor * sin = ggml_sin(ctx, input); + REQUIRE(sin != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, sin); + + 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())); + } + + ggml_free(ctx); +} + +static void run_zero_output_device_support_checks() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 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, 3840, 0); + ggml_tensor * rhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3840, 0); + REQUIRE(lhs != nullptr); + REQUIRE(rhs != nullptr); + + ggml_tensor * sum = ggml_add(ctx, lhs, rhs); + ggml_tensor * prod = ggml_mul(ctx, lhs, rhs); + ggml_tensor * scaled = ggml_scale(ctx, lhs, 2.0f); + REQUIRE(sum != nullptr); + REQUIRE(prod != nullptr); + REQUIRE(scaled != nullptr); + REQUIRE(ggml_backend_supports_op(backend, sum)); + REQUIRE(ggml_backend_supports_op(backend, prod)); + REQUIRE(ggml_backend_supports_op(backend, scaled)); + + ggml_tensor * acc = ggml_acc(ctx, lhs, rhs, lhs->nb[1], lhs->nb[2], lhs->nb[3], 0); + ggml_tensor * set = ggml_set(ctx, lhs, rhs, lhs->nb[1], lhs->nb[2], lhs->nb[3], 0); + ggml_tensor * cpy = ggml_cpy(ctx, lhs, rhs); + REQUIRE(acc != nullptr); + REQUIRE(set != nullptr); + REQUIRE(cpy != nullptr); + REQUIRE(!ggml_backend_supports_op(backend, acc)); + REQUIRE(!ggml_backend_supports_op(backend, set)); + REQUIRE(ggml_backend_supports_op(backend, cpy)); + + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_scale_f32_device_support_checks() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 32); + ggml_tensor * out = ggml_scale_bias(ctx, input, 1.5f, 0.5f); + ggml_tensor * inplace = ggml_scale_bias_inplace(ctx, input, -2.0f, 1.0f); + ggml_tensor * view = ggml_view_1d(ctx, input, 16, 8 * sizeof(float)); + ggml_tensor * view_out = ggml_scale_bias(ctx, view, 0.25f, -0.5f); + ggml_tensor * view_inpl = ggml_scale_bias_inplace(ctx, view, 0.5f, 0.25f); + REQUIRE(input != nullptr); + REQUIRE(out != nullptr); + REQUIRE(inplace != nullptr); + REQUIRE(view != nullptr); + REQUIRE(view_out != nullptr); + REQUIRE(view_inpl != nullptr); + + REQUIRE(ggml_backend_supports_op(backend, out)); + REQUIRE(ggml_backend_supports_op(backend, inplace)); + REQUIRE(ggml_backend_supports_op(backend, view_out)); + REQUIRE(ggml_backend_supports_op(backend, view_inpl)); + + ggml_free(ctx); + ggml_backend_free(backend); +} + static void run_transient_import_checks() { ggml_init_params params = {}; params.mem_size = 256 * 1024; @@ -3757,7 +11570,7 @@ static void run_transient_import_checks() { ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); ggml_tensor * sum = ggml_add(ctx, a, b); - ggml_tensor * out = ggml_sqr(ctx, sum); + ggml_tensor * out = ggml_sin(ctx, sum); REQUIRE(a != nullptr); REQUIRE(b != nullptr); REQUIRE(sum != nullptr); @@ -3771,7 +11584,7 @@ static void run_transient_import_checks() { REQUIRE(imported.valid()); REQUIRE(imported.graph.nodes().size() == 2); REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); - REQUIRE(imported.graph.nodes()[1].op == GGML_OP_SQR); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_SIN); const ggml::hrx::Value * a_value = imported.graph.values().find_tensor(a); const ggml::hrx::Value * sum_value = imported.graph.values().find_tensor(sum); @@ -3794,6 +11607,68 @@ static void run_transient_import_checks() { ggml_free(ctx); } +static void run_graph_view_preserves_shared_output_storage_checks() { + ggml_init_params params = {}; + params.mem_size = 256 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * out0 = ggml_add(ctx, sum, c); + ggml_tensor * out1 = ggml_add(ctx, sum, d); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(sum != nullptr); + REQUIRE(out0 != nullptr); + REQUIRE(out1 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + ggml_build_forward_expand(graph, out1); + REQUIRE(graph->n_nodes == 3); + REQUIRE(graph->nodes[0] == sum); + REQUIRE(graph->nodes[1] == out0); + REQUIRE(graph->nodes[2] == out1); + REQUIRE(ggml_node_get_use_count(graph, 0) == 2); + + ggml_cgraph graph_view = ggml_graph_view(graph, 0, 2); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(graph_view); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_ADD); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_ADD); + + const ggml::hrx::Value * sum_value = imported.graph.values().find_tensor(sum); + REQUIRE(sum_value != nullptr); + REQUIRE(sum_value->kind == ggml::hrx::ValueKind::External); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 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(commands.commands.size() == 2); + REQUIRE(commands.commands[0].bindings[2].value == sum_value->id); + REQUIRE(commands.commands[0].bindings[2].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(commands.commands[1].bindings[0].value == sum_value->id); + REQUIRE(commands.commands[1].bindings[0].origin == ggml::hrx::CommandBindingOrigin::GraphValue); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, sum_value->id) == nullptr); + REQUIRE(command_program_verifies(commands)); + + ggml_free(ctx); +} + static void run_chained_dispatch_requires_transients() { ggml_init_params params = {}; params.mem_size = 256 * 1024; @@ -4149,6 +12024,20 @@ static void run_disjoint_transient_plan_packing_checks() { ggml_free(ctx); } +static void run_command_shape_hash_checks() { + REQUIRE(ggml::hrx::command_program_shape_hash("") == UINT64_C(1469598103934665603)); + std::string bytes; + for (int i = 0; i < 256; ++i) { + bytes.push_back(static_cast(i)); + } + bytes.append("\0key\xff", 5); + REQUIRE(ggml::hrx::command_program_shape_hash(bytes) == UINT64_C(6643458358697029213)); + ggml::hrx::GraphProgram program(1, "gfx1151", nullptr, nullptr, bytes); + bytes.clear(); + REQUIRE(program.command_shape_hash() == UINT64_C(6643458358697029213)); + REQUIRE(program.command_shape_hash() == ggml::hrx::command_program_shape_hash(program.command_shape())); +} + static void run_graph_program_cache_uid_mismatch_checks() { ggml::hrx::GraphProgramCache cache; const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); @@ -4179,6 +12068,8 @@ static void run_graph_program_cache_uid_mismatch_checks() { ggml::hrx::GraphProgramLookup lookup = cache.get_or_build(*graph0, corpus, "gfx1151"); REQUIRE(lookup.valid()); + const uint64_t first_shape_hash = lookup.program->command_shape_hash(); + REQUIRE(first_shape_hash == ggml::hrx::command_program_shape_hash(lookup.program->command_shape())); ggml::hrx::GraphProgramCacheStats stats = cache.stats(); REQUIRE(stats.builds == 1); REQUIRE(stats.hits == 0); @@ -4188,6 +12079,7 @@ static void run_graph_program_cache_uid_mismatch_checks() { stats = cache.stats(); REQUIRE(stats.builds == 1); REQUIRE(stats.hits == 1); + REQUIRE(lookup.program->command_shape_hash() == first_shape_hash); ggml_cgraph * graph1 = ggml_new_graph(ctx); REQUIRE(graph1 != nullptr); @@ -4200,8 +12092,11 @@ static void run_graph_program_cache_uid_mismatch_checks() { stats = cache.stats(); REQUIRE(stats.builds == 2); REQUIRE(stats.hits == 1); + REQUIRE(lookup.program->command_shape_hash() != first_shape_hash); + REQUIRE(lookup.program->command_shape_hash() == + ggml::hrx::command_program_shape_hash(lookup.program->command_shape())); - ggml_tensor * unsupported = ggml_sqr(ctx, a); + ggml_tensor * unsupported = ggml_sin(ctx, a); REQUIRE(unsupported != nullptr); ggml_cgraph * graph2 = ggml_new_graph(ctx); REQUIRE(graph2 != nullptr); @@ -4245,6 +12140,153 @@ static void run_graph_program_cache_uid_mismatch_checks() { ggml_free(ctx); } +static void run_graph_match_hash_collision_checks() { + ggml_context * ctx = ggml_init({ 2 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + const size_t buckets = ggml_hash_size(10); + std::vector seen(buckets, nullptr); + ggml_tensor * a = nullptr; + ggml_tensor * b = nullptr; + for (size_t i = 0; i <= buckets && b == nullptr; ++i) { + ggml_tensor * tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + const size_t bucket = ggml_hash(tensor) % buckets; + if (seen[bucket] != nullptr) { + a = seen[bucket]; + b = tensor; + } else { + seen[bucket] = tensor; + } + } + REQUIRE(a != nullptr && b != nullptr && a != b); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * out = ggml_add(ctx, sum, c); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 256, false); + ggml_build_forward_expand(graph, out); + ggml::hrx::GraphProgramCache cache; + auto lookup = cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.program->graph().values().size() == 5); + const auto match = lookup.program->match_current_graph(*graph); + REQUIRE(match.valid()); + for (const ggml_tensor * tensor : { a, b, c, out }) { + REQUIRE(std::any_of(match.external_bindings.begin(), match.external_bindings.end(), + [&](const auto & binding) { return binding.tensor == tensor; })); + } + for (int i = 0; i < 128; ++i) { + out = ggml_add(ctx, out, i % 2 == 0 ? a : b); + } + ggml_build_forward_expand(graph, out); + lookup = cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.program->graph().values().size() == 133); + REQUIRE(lookup.program->match_current_graph(*graph).valid()); + ggml_free(ctx); +} + +static void run_graph_match_bijection_checks() { + ggml_context * ctx = ggml_init({ 512 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * sum = ggml_add(ctx, a, b); + ggml_tensor * out = ggml_add(ctx, sum, b); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, out); + ggml::hrx::GraphProgramCache cache; + auto lookup = cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(lookup.valid()); + REQUIRE(lookup.program->match_current_graph(*graph).valid()); + + sum->src[1] = a; + auto duplicate_tensor = lookup.program->match_current_graph(*graph); + REQUIRE(!duplicate_tensor.valid()); + REQUIRE(status_contains(duplicate_tensor.status, "tensor maps to cached values")); + sum->src[1] = b; + out->src[1] = c; + auto duplicate_value = lookup.program->match_current_graph(*graph); + REQUIRE(!duplicate_value.valid()); + REQUIRE(status_contains(duplicate_value.status, "maps to multiple current tensors")); + out->src[1] = b; + REQUIRE(lookup.program->match_current_graph(*graph).valid()); + ggml_free(ctx); +} + +static void run_validated_graph_uid_match_checks() { + ggml::hrx::GraphProgramCache cache; + const auto & corpus = ggml::hrx::get_qwen_kernel_corpus(); + ggml_context * ctx = ggml_init({ 2 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + auto make_graph = [&](uint64_t uid, int64_t elements, bool two_nodes) { + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, elements); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, elements); + ggml_tensor * out = ggml_add(ctx, a, b); + if (two_nodes) { + out = ggml_add(ctx, out, b); + } + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, out); + graph->uid = uid; + return std::make_pair(graph, a); + }; + auto first = make_graph(4001, 8, false); + auto alias = make_graph(4002, 8, false); + auto original = cache.get_or_build(*first.first, corpus, "gfx1151"); + REQUIRE(original.valid()); + for (int i = 0; i < 2; ++i) { + auto match = cache.get_or_build(*alias.first, corpus, "gfx1151"); + REQUIRE(match.valid()); + REQUIRE(match.program == original.program); + REQUIRE(std::any_of(match.match.external_bindings.begin(), match.match.external_bindings.end(), + [&](const auto & binding) { return binding.tensor == alias.second; })); + } + REQUIRE(cache.stats().builds == 1); + + auto new_nodes = make_graph(4002, 16, false); + auto changed = cache.get_or_build(*new_nodes.first, corpus, "gfx1151"); + REQUIRE(changed.valid()); + REQUIRE(changed.program->uid() == 4002); + REQUIRE(cache.stats().builds == 2); + + auto second_alias = make_graph(4003, 8, false); + REQUIRE(cache.get_or_build(*second_alias.first, corpus, "gfx1151").program == original.program); + auto replacement = make_graph(4001, 8, true); + REQUIRE(cache.get_or_build(*replacement.first, corpus, "gfx1151").valid()); + REQUIRE(cache.stats().builds == 3); + auto after_replacement = cache.get_or_build(*second_alias.first, corpus, "gfx1151"); + REQUIRE(after_replacement.valid()); + REQUIRE(after_replacement.program->uid() == 4003); + REQUIRE(cache.stats().builds == 4); + + auto checked_alias = make_graph(4004, 8, false); + REQUIRE(cache.get_or_build(*checked_alias.first, corpus, "gfx1151").valid()); + const char * validate_name = "GGML_HRX_VALIDATE_GRAPH_UID_CACHE"; + const char * validate_value = std::getenv(validate_name); + const bool had_validate = validate_value != nullptr; + const std::string saved_validate = had_validate ? validate_value : ""; + REQUIRE(setenv(validate_name, "1", 1) == 0); + checked_alias.second->ne[0] = 12; + REQUIRE(!cache.get_or_build(*checked_alias.first, corpus, "gfx1151").valid()); + checked_alias.second->ne[0] = 8; + REQUIRE(cache.get_or_build(*checked_alias.first, corpus, "gfx1151").valid()); + restore_environment_value(validate_name, had_validate, saved_validate); + + for (uint64_t uid = 4100; uid < 4230; ++uid) { + checked_alias.first->uid = uid; + auto match = cache.get_or_build(*checked_alias.first, corpus, "gfx1151"); + REQUIRE(match.valid()); + REQUIRE(match.program->uid() == 4003); + } + REQUIRE(cache.stats().builds == 4); + cache.clear(); + auto after_clear = cache.get_or_build(*second_alias.first, corpus, "gfx1151"); + REQUIRE(after_clear.valid()); + REQUIRE(after_clear.program->uid() == 4003); + REQUIRE(cache.stats().builds == 5); + ggml_free(ctx); +} + static void run_graph_executor_contract_checks() { ggml_backend_hrx_device_context device_context = {}; ggml_backend_hrx_context backend_context = {}; @@ -4261,11 +12303,11 @@ static void run_graph_executor_contract_checks() { ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); ggml_tensor * add_out = ggml_add(ctx, a, b); - ggml_tensor * sqr_out = ggml_sqr(ctx, a); + ggml_tensor * sin_out = ggml_sin(ctx, a); REQUIRE(a != nullptr); REQUIRE(b != nullptr); REQUIRE(add_out != nullptr); - REQUIRE(sqr_out != nullptr); + REQUIRE(sin_out != nullptr); ggml_cgraph * add_graph = ggml_new_graph(ctx); REQUIRE(add_graph != nullptr); @@ -4274,21 +12316,115 @@ static void run_graph_executor_contract_checks() { REQUIRE(add_support.supported); REQUIRE(add_support.status.success()); - const ggml::hrx::GraphExecutionResult missing_binding = executor.execute(*add_graph); - REQUIRE(!missing_binding.success()); - REQUIRE(missing_binding.code == GGML_STATUS_FAILED); - REQUIRE(status_contains(missing_binding.status, "external value")); - REQUIRE(status_contains(missing_binding.status, "not bound")); + const ggml::hrx::GraphExecutionResult missing_binding = executor.execute(*add_graph); + REQUIRE(!missing_binding.success()); + REQUIRE(missing_binding.code == GGML_STATUS_FAILED); + REQUIRE(status_contains(missing_binding.status, "external value")); + REQUIRE(status_contains(missing_binding.status, "not bound")); + + ggml_cgraph * sin_graph = ggml_new_graph(ctx); + REQUIRE(sin_graph != nullptr); + ggml_build_forward_expand(sin_graph, sin_out); + const ggml::hrx::GraphSupportResult sin_support = executor.can_execute(*sin_graph); + REQUIRE(!sin_support.supported); + REQUIRE(status_contains(sin_support.status, "unsupported HRX node 0")); + REQUIRE(status_contains(sin_support.status, "SIN")); + + ggml_free(ctx); +} + +static void build_add_graph_program(ggml::hrx::GraphProgramCache & cache, uint64_t uid, bool two_outputs) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out0 = ggml_add(ctx, a, b); + REQUIRE(a != nullptr); + REQUIRE(b != nullptr); + REQUIRE(out0 != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out0); + if (two_outputs) { + ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * d = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); + ggml_tensor * out1 = ggml_add(ctx, c, d); + REQUIRE(c != nullptr); + REQUIRE(d != nullptr); + REQUIRE(out1 != nullptr); + ggml_build_forward_expand(graph, out1); + } + graph->uid = uid; + + const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); + ggml::hrx::GraphProgramLookup lookup = cache.get_or_build(*graph, corpus, "gfx1151"); + REQUIRE(lookup.valid()); + ggml_free(ctx); +} + +static void run_command_program_kernel_dump_checks() { + static constexpr const char * kDumpEnv = "GGML_HRX_DUMP_COMMAND_PROGRAM_DIR"; + EnvironmentVariableGuard env(kDumpEnv); - ggml_cgraph * sqr_graph = ggml_new_graph(ctx); - REQUIRE(sqr_graph != nullptr); - ggml_build_forward_expand(sqr_graph, sqr_out); - const ggml::hrx::GraphSupportResult sqr_support = executor.can_execute(*sqr_graph); - REQUIRE(!sqr_support.supported); - REQUIRE(status_contains(sqr_support.status, "unsupported HRX node 0")); - REQUIRE(status_contains(sqr_support.status, "SQR")); + { + const std::filesystem::path dump_dir = fresh_test_directory("kernel-dump-disabled"); + env.unset(); + ggml::hrx::GraphProgramCache cache; + build_add_graph_program(cache, 5001, false); + REQUIRE(list_directories(dump_dir).empty()); + std::filesystem::remove_all(dump_dir); + } - ggml_free(ctx); + const std::filesystem::path dump_dir = fresh_test_directory("kernel-dump-enabled"); + env.set(dump_dir); + + ggml::hrx::GraphProgramCache cache; + build_add_graph_program(cache, 5002, false); + std::vector dumps = list_directories(dump_dir); + REQUIRE(dumps.size() == 1); + REQUIRE(std::filesystem::exists(dumps[0] / "kernels.txt")); + REQUIRE(std::filesystem::exists(dumps[0] / "kernels.dot")); + + const std::string add_kernels = read_text_file(dumps[0] / "kernels.txt"); + const std::string add_dot = read_text_file(dumps[0] / "kernels.dot"); + REQUIRE(string_contains(add_kernels, "main 0 loom_libs:ggml_binary_f32")); + REQUIRE(!string_contains(add_kernels, "binding")); + REQUIRE(!string_contains(add_kernels, "transient")); + REQUIRE(string_contains(add_dot, "digraph hrx_kernel_invocations")); + REQUIRE(string_contains(add_dot, "main 0\\nloom_libs:ggml_binary_f32")); + + build_add_graph_program(cache, 5003, false); + dumps = list_directories(dump_dir); + REQUIRE(dumps.size() == 1); + + const std::filesystem::path second_dump_dir = fresh_test_directory("kernel-dump-second-directory"); + env.set(second_dump_dir); + build_add_graph_program(cache, 5005, false); + std::vector second_dumps = list_directories(second_dump_dir); + REQUIRE(second_dumps.size() == 1); + REQUIRE(std::filesystem::exists(second_dumps[0] / "kernels.txt")); + std::filesystem::remove_all(second_dump_dir); + env.set(dump_dir); + + build_add_graph_program(cache, 5004, true); + dumps = list_directories(dump_dir); + REQUIRE(dumps.size() == 2); + + const std::string two_add_kernels = read_text_file(dumps[1] / "kernels.txt"); + const std::string two_add_dot = read_text_file(dumps[1] / "kernels.dot"); + REQUIRE(string_contains(two_add_kernels, "main 0 loom_libs:ggml_binary_f32")); + REQUIRE(string_contains(two_add_kernels, "main 1 loom_libs:ggml_binary_f32 deps=0")); + REQUIRE(string_contains(two_add_dot, "main0 -> main1")); + REQUIRE(!std::filesystem::exists(dumps[1] / "commands.txt")); + REQUIRE(!std::filesystem::exists(dumps[1] / "commands.json")); + REQUIRE(!std::filesystem::exists(dumps[1] / "manifest.json")); + + std::filesystem::remove_all(dump_dir); } static std::vector read_transient_i32(ggml_backend_hrx_context * context, @@ -4314,6 +12450,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); @@ -4485,6 +12829,250 @@ static void run_add_f32() { ggml_backend_free(backend); } +static std::vector run_rope_scale_numerical_case(ggml_backend_t backend, bool per_token, int mode) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t head_size = 128; + constexpr int64_t n_dims = 96; + constexpr int64_t head_count = 2; + constexpr int64_t token_count = 3; + constexpr int64_t elements = head_size * head_count * token_count; + + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); + ggml_tensor * rope = + ggml_rope_ext(ctx, input, pos, freqs, n_dims, mode, 0, 10000.0f, 1.0f, 0.0f, 1.19024f, 32.0f, 1.0f); + ggml_tensor * scales = per_token ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count) : nullptr; + ggml_tensor * output = per_token ? ggml_mul(ctx, rope, scales) : ggml_scale(ctx, rope, 0.0883883461f); + REQUIRE(input != nullptr); + REQUIRE(pos != nullptr); + REQUIRE(freqs != nullptr); + REQUIRE(rope != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + std::vector input_data(elements); + std::vector freq_data(n_dims / 2); + const std::vector positions = { 3, 7, 11 }; + const std::vector scale_data = { 0.125f, -0.25f, 0.5f }; + for (int64_t i = 0; i < elements; ++i) { + input_data[i] = static_cast((i * 17) % 113) * 0.03125f - 1.5f; + } + for (int64_t i = 0; i < n_dims / 2; ++i) { + freq_data[i] = 0.75f + static_cast(i % 7) * 0.125f; + } + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + ggml_backend_tensor_set(pos, positions.data(), 0, positions.size() * sizeof(int32_t)); + ggml_backend_tensor_set(freqs, freq_data.data(), 0, freq_data.size() * sizeof(float)); + if (per_token) { + ggml_backend_tensor_set(scales, scale_data.data(), 0, scale_data.size() * sizeof(float)); + } + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + std::vector actual(elements); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return actual; +} + +static void run_rope_scale_f32_numerics() { + ggml_backend_t hrx = ggml_backend_hrx_init(0); + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(hrx != nullptr); + REQUIRE(cpu != nullptr); + + for (const bool per_token : { false, true }) { + for (const int mode : { static_cast(GGML_ROPE_TYPE_NORMAL), static_cast(GGML_ROPE_TYPE_NEOX) }) { + const std::vector expected = run_rope_scale_numerical_case(cpu, per_token, mode); + const std::vector actual = run_rope_scale_numerical_case(hrx, per_token, mode); + REQUIRE(actual.size() == expected.size()); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = 2.0e-4f + 2.0e-4f * std::fabs(expected[i]); + REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); + } + } + } + + ggml_backend_free(cpu); + ggml_backend_free(hrx); +} + +struct GatherAddRmsNormOutputs { + std::vector raw; + std::vector normalized; +}; + +static GatherAddRmsNormOutputs run_gather_add_rmsnorm_numerical_case(ggml_backend_t backend) { + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t hidden_size = 128; + constexpr int64_t source_token_count = 4; + constexpr int64_t output_token_count = 5; + GatherAddRmsNormGraph tensors = + build_gather_add_rmsnorm_graph(ctx, hidden_size, source_token_count, output_token_count); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector attention(hidden_size * source_token_count); + std::vector residual(hidden_size * source_token_count); + std::vector weight(hidden_size); + const std::vector row_ids = { 3, 1, 3, 0, 2 }; + for (size_t i = 0; i < attention.size(); ++i) { + attention[i] = static_cast((i * 13) % 47) * 0.03125f - 0.5f; + residual[i] = static_cast((i * 7) % 31) * -0.015625f + 0.25f; + } + for (size_t i = 0; i < weight.size(); ++i) { + weight[i] = 0.5f + static_cast(i % 11) * 0.0625f; + } + ggml_backend_tensor_set(tensors.attention, attention.data(), 0, attention.size() * sizeof(float)); + ggml_backend_tensor_set(tensors.residual, residual.data(), 0, residual.size() * sizeof(float)); + ggml_backend_tensor_set(tensors.row_ids, row_ids.data(), 0, row_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(tensors.weight, weight.data(), 0, weight.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, tensors.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + GatherAddRmsNormOutputs outputs; + outputs.raw.resize(hidden_size * output_token_count); + outputs.normalized.resize(hidden_size * output_token_count); + ggml_backend_tensor_get(tensors.raw_output, outputs.raw.data(), 0, outputs.raw.size() * sizeof(float)); + ggml_backend_tensor_get(tensors.normalized_output, outputs.normalized.data(), 0, + outputs.normalized.size() * sizeof(float)); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return outputs; +} + +static void run_gather_add_rmsnorm_f32_numerics() { + ggml_backend_t hrx = ggml_backend_hrx_init(0); + ggml_backend_t cpu = ggml_backend_cpu_init(); + REQUIRE(hrx != nullptr); + REQUIRE(cpu != nullptr); + + const GatherAddRmsNormOutputs expected = run_gather_add_rmsnorm_numerical_case(cpu); + const GatherAddRmsNormOutputs actual = run_gather_add_rmsnorm_numerical_case(hrx); + REQUIRE(actual.raw.size() == expected.raw.size()); + REQUIRE(actual.normalized.size() == expected.normalized.size()); + for (size_t i = 0; i < actual.raw.size(); ++i) { + REQUIRE(actual.raw[i] == expected.raw[i]); + const float tolerance = 3.0e-4f + 3.0e-4f * std::fabs(expected.normalized[i]); + REQUIRE(std::fabs(actual.normalized[i] - expected.normalized[i]) <= tolerance); + } + + ggml_backend_free(cpu); + ggml_backend_free(hrx); +} + +static void run_scale_f32() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t element_count = 1024; + ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * out = ggml_scale_bias(ctx, input, -1.5f, 0.25f); + REQUIRE(input != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector input_data(element_count); + std::vector expected(element_count); + for (int64_t i = 0; i < element_count; ++i) { + input_data[i] = static_cast(i % 23) * 0.125f - 1.0f; + expected[i] = input_data[i] * -1.5f + 0.25f; + } + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(element_count); + ggml_backend_tensor_get(out, actual.data(), 0, actual.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(std::fabs(actual[i] - expected[i]) <= 1.0e-6f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_scale_f32_inplace() { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 512 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t element_count = 1024; + ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, element_count); + ggml_tensor * out = ggml_scale_bias_inplace(ctx, input, -1.5f, 0.25f); + REQUIRE(input != nullptr); + REQUIRE(out != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, out); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + std::vector input_data(element_count); + std::vector expected(element_count); + for (int64_t i = 0; i < element_count; ++i) { + input_data[i] = static_cast(i % 23) * 0.125f - 1.0f; + expected[i] = input_data[i] * -1.5f + 0.25f; + } + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(element_count); + ggml_backend_tensor_get(out, actual.data(), 0, actual.size() * sizeof(float)); + for (int64_t i = 0; i < element_count; ++i) { + REQUIRE(std::fabs(actual[i] - expected[i]) <= 1.0e-6f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + static void run_two_independent_add_f32() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -4683,7 +13271,7 @@ static void run_chained_add_f32() { ggml_backend_free(backend); } -static void run_same_uid_distinct_graph_reuses_graph_program() { +static void run_same_uid_distinct_graph_reuses_graph_program(bool distinct_uid = false) { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -4746,7 +13334,7 @@ static void run_same_uid_distinct_graph_reuses_graph_program() { ggml_cgraph * graph1 = ggml_new_graph(ctx1); REQUIRE(graph1 != nullptr); ggml_build_forward_expand(graph1, out1); - graph1->uid = 1003; + graph1->uid = distinct_uid ? 1004 : 1003; ggml_backend_buffer_t buffer1 = ggml_backend_alloc_ctx_tensors(ctx1, backend); REQUIRE(buffer1 != nullptr); @@ -4773,6 +13361,23 @@ static void run_same_uid_distinct_graph_reuses_graph_program() { REQUIRE(actual[i] == expected[i]); } + if (distinct_uid) { + for (int64_t i = 0; i < element_count; ++i) { + a_data[i] += 1.0f; + expected[i] = a_data[i] + b_data[i]; + } + ggml_backend_tensor_set(a1, a_data.data(), 0, a_data.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, graph1) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + ggml_backend_tensor_get(out1, actual.data(), 0, actual.size() * sizeof(float)); + REQUIRE(actual == expected); + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &cache_stats)); + REQUIRE(cache_stats.graph_program_builds == 1); + REQUIRE(cache_stats.graph_program_hits == 2); + REQUIRE(cache_stats.prepared_program_builds == 2); + REQUIRE(cache_stats.prepared_program_hits == 1); + } + ggml_backend_buffer_free(buffer1); ggml_free(ctx1); ggml_backend_buffer_free(buffer0); @@ -4791,7 +13396,7 @@ static void run_unsupported_op_fails() { REQUIRE(ctx != nullptr); ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 8); - ggml_tensor * out = ggml_sqr(ctx, a); + ggml_tensor * out = ggml_sin(ctx, a); REQUIRE(a != nullptr); REQUIRE(out != nullptr); @@ -4811,47 +13416,464 @@ static void run_unsupported_op_fails() { ggml_backend_free(backend); } -int main() { - run_status_checks(); - run_command_plan_metadata_checks(); - run_dispatch_registry_checks(); - run_graph_import_checks(); - run_graph_snapshot_diagnostics_checks(); - run_unmatched_graph_diagnostics_checks(); - run_completion_counter_plan_checks(); - run_graph_index_checks(); - run_graph_traversal_checks(); - run_qwen_token_embedding_dispatch_checks(); - run_gather_add_dispatch_checks(); - run_qwen_flash_attention_dispatch_checks(); - run_qwen_attention_postprocess_dispatch_checks(); - run_qwen_matmul_dispatch_checks(); - schedule_qwen_terminal_q6k_q8_command(1); - schedule_qwen_terminal_q6k_q8_command(18); - run_qwen_router_top8_dispatch_checks(); - run_qwen_routed_gate_up_dispatch_checks(); - run_alias_value_import_checks(); - run_multi_dispatch_checks(); - run_layout_alias_scheduler_elision_checks(); - run_transient_import_checks(); - run_chained_dispatch_requires_transients(); - run_graph_replay_host_staging_is_not_ineligible(); - run_multiple_transient_plan_checks(); - run_disjoint_transient_plan_packing_checks(); - run_graph_program_cache_uid_mismatch_checks(); - run_graph_executor_contract_checks(); - REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); - - if (ggml_backend_hrx_get_device_count() == 0) { - std::fprintf(stderr, "test skipped: no HRX devices available\n"); - return 0; - } - - run_qwen_expert_table_partition_prefill_512_execution(); - run_add_f32(); - run_two_independent_add_f32(); - run_chained_add_f32(); - run_same_uid_distinct_graph_reuses_graph_program(); - run_unsupported_op_fails(); - return 0; +static void run_lowtoken_residual_dispatch_checks() { + struct Case { + ggml_type type; + int64_t k, n, tokens; + int layout; + int side_use; + bool alias_weight; + bool second_add; + bool fused; + }; + + const Case cases[] = { + { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 17408, 5120, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q6_K, 17408, 5120, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 4096, 4096, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 32768, 4096, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 1, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 2, 1, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 3, 1, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 4, 1, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 0, false, false, true }, + { GGML_TYPE_Q6_K, 17408, 5120, 5, 1, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, false, true, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 0, false, true, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 1, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 1, 1, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 1, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 1, 2, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 2, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 2, 0, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 2, 0, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 3, 0, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 3, 0, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, true, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 0, true, false, false }, + { GGML_TYPE_Q4_K, 3840, 4096, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 4032, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 4096, 8192, 1, 0, 0, false, false, true }, + { GGML_TYPE_F16, 6144, 5120, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 6, 1, 0, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 5120, 256, 1, 0, false, false, false }, + }; + for (const Case & c : cases) { + ggml_context * ctx = ggml_init({ 16 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, c.type, c.k, c.n + (c.alias_weight ? 64 : 0)); + if (c.alias_weight) { + weight = ggml_view_2d(ctx, weight, c.k, c.n, weight->nb[1], 0); + } + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c.k, c.tokens); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * laid_out = projection; + if (c.layout == 1) { + laid_out = ggml_reshape_2d(ctx, projection, c.n, c.tokens); + } else if (c.layout == 2) { + laid_out = ggml_reshape_2d(ctx, projection, c.n / 2, c.tokens * 2); + } else if (c.layout == 3) { + laid_out = ggml_view_2d(ctx, projection, c.n - 64, c.tokens, projection->nb[1], 0); + } + ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, laid_out->ne[0], laid_out->ne[1]); + ggml_tensor * output = ggml_add(ctx, laid_out, residual); + if (c.second_add) { + output = ggml_add(ctx, output, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, laid_out->ne[0], laid_out->ne[1])); + } + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + if (c.side_use != 0) { + ggml_build_forward_expand(graph, ggml_scale(ctx, c.side_use == 1 ? projection : laid_out, 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::CommandProgram program = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(program.valid()); + REQUIRE(command_program_verifies(program)); + size_t fused = 0, binary = 0; + bool fused_second_add = false; + for (const auto & command : program.commands) { + const std::string name = kernel_name_for_id(command.kernel.kernel_id); + binary += name == "loom_libs:ggml_binary_f32"; + const bool tiled_postops = name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; + const bool skinny_postops = name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; + const bool vector_postops = name == "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32"; + if (!tiled_postops && !skinny_postops && !vector_postops) { + continue; + } + ++fused; + if (vector_postops) { + fused_second_add = + fused_second_add || + (command.kernel.compile_parameters.at("ggml.matmul.vector_postops.apply_bias") == "1" && + command.kernel.compile_parameters.at("ggml.matmul.vector_postops.apply_residual") == "1"); + } else { + fused_second_add = + fused_second_add || command.kernel.compile_parameters.at("ggml.mul_mat_postops.epilogue") == "3"; + } + REQUIRE(command.bindings.size() >= 4); + REQUIRE(command.kernel.integer_parameters.at("token_count") == c.tokens); + if (vector_postops) { + REQUIRE(command.kernel.compile_parameters.count("ggml.matmul.vector_postops.weight_format") == 1); + } else { + REQUIRE(command.kernel.compile_parameters.count("ggml.mul_mat_postops.weight_format") == 1); + } + } + REQUIRE(fused == (c.fused ? 1 : 0)); + const size_t expected_binary = c.fused ? (c.second_add && !fused_second_add ? 1 : 0) : 1; + REQUIRE(binary == expected_binary); + ggml_free(ctx); + } +} + +static void run_lowtoken_bias_dispatch_checks() { + ggml_context * ctx = ggml_init({ 16 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 3584, 1024); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3584, 2); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1024); + ggml_tensor * output = ggml_add(ctx, projection, bias); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + const ggml::hrx::CommandProgram program = ggml::hrx::build_command_program( + imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(program.valid()); + REQUIRE(command_program_verifies(program)); + + size_t skinny_bias = 0; + size_t tiled_bias = 0; + for (const auto & command : program.commands) { + const std::string name = kernel_name_for_id(command.kernel.kernel_id); + skinny_bias += name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; + tiled_bias += name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; + } + REQUIRE(skinny_bias == 1); + REQUIRE(tiled_bias == 0); + + ggml_free(ctx); +} + +static void run_gdn_selected_rms_q8_dispatch_checks() { + ggml_init_params params = {}; + params.mem_size = 32 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + enum { + UnreadyInput = 1, + UnreadyGate = 2, + GatherSide = 4, + AttentionSide = 8, + GdnObserved = 16, + GatherObserved = 32, + Unaligned = 64, + CacheRead = 128, + OutputSide = 256, + CachedGate = 512, + TanhGate = 1024, + AttentionReshape = 2048, + AttentionObserved = 4096, + }; + + const struct { + int64_t tokens, heads, sequences; + int flags; + ggml_type type; + bool fused; + } cases[] = { + {1,48,1,0,GGML_TYPE_Q4_K,true}, + {2,48,1,0,GGML_TYPE_Q4_K,true}, + {3,48,1,0,GGML_TYPE_Q4_K,true}, + {4,48,1,0,GGML_TYPE_Q4_K,true}, + {5,48,1,0,GGML_TYPE_Q4_K,true}, + {1,24,1,0,GGML_TYPE_Q4_K,true}, + {5,24,1,0,GGML_TYPE_Q4_K,true}, + {1,48,1,0,GGML_TYPE_Q6_K,true}, + {5,48,1,0,GGML_TYPE_Q6_K,true}, + {1,48,1,AttentionReshape,GGML_TYPE_Q4_K,true}, + {5,48,1,AttentionReshape,GGML_TYPE_Q4_K,true}, + {1,48,1,UnreadyInput,GGML_TYPE_Q4_K,false}, + {1,48,1,UnreadyGate,GGML_TYPE_Q4_K,false}, + {5,48,1,UnreadyGate,GGML_TYPE_Q4_K,false}, + {1,48,1,GatherSide,GGML_TYPE_Q4_K,false}, + {5,48,1,AttentionSide,GGML_TYPE_Q4_K,false}, + {1,48,1,GdnObserved,GGML_TYPE_Q4_K,false}, + {1,48,1,GatherObserved,GGML_TYPE_Q4_K,false}, + {1,48,1,AttentionObserved,GGML_TYPE_Q4_K,false}, + {1,48,1,Unaligned,GGML_TYPE_Q4_K,false}, + {5,48,1,CacheRead,GGML_TYPE_Q4_K,false}, + {1,48,1,OutputSide,GGML_TYPE_Q4_K,false}, + {1,48,1,CachedGate,GGML_TYPE_Q4_K,false}, + {5,48,1,TanhGate,GGML_TYPE_Q4_K,false}, + {1,48,1,0,GGML_TYPE_F16,false}, + {6,48,1,0,GGML_TYPE_Q4_K,false}, + {1,48,2,0,GGML_TYPE_Q4_K,false}, + {1,6,1,0,GGML_TYPE_Q4_K,false}, + }; + for (const auto & test : cases) { + const int64_t width = 128, heads = test.heads, qheads = heads / 3; + const int64_t hidden = width * (heads + 2 * qheads), channels = width * heads; + const int64_t state_elements = width * channels; + const int64_t written = std::min(test.tokens, int64_t{5}); + const size_t state_bytes = state_elements * sizeof(float); + ggml_tensor * raw = ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, test.tokens * test.sequences), 0.5f); + ggml_tensor * alpha = ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, heads, test.tokens * test.sequences), 0.5f); + ggml_tensor * beta_raw = ggml_scale(ctx, ggml_dup_tensor(ctx, alpha), 0.5f); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, heads); + ggml_tensor * scale = ggml_dup_tensor(ctx, bias); + ggml_tensor * gate = ggml_mul(ctx, ggml_softplus(ctx, ggml_add(ctx, ggml_reshape_3d(ctx, alpha, heads, test.tokens, test.sequences), bias)), scale); + gate = ggml_reshape_4d(ctx, gate, 1, heads, test.tokens, test.sequences); + ggml_tensor * beta = ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta_raw, 1, heads, test.tokens, test.sequences)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 20); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, test.sequences); + ggml_tensor * gathered = ggml_get_rows(ctx, cache, ids); + if (test.flags & GatherObserved) { + ggml_set_output(gathered); + } + ggml_tensor * state = ggml_reshape_4d(ctx, gathered, width, width, heads, test.sequences); + const auto view = [&](int64_t nheads, size_t offset) { + return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, width * sizeof(float), + hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); + }; + ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); + ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); + ggml_tensor * v = view(heads, 2 * width * qheads * sizeof(float)); + ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, gate, beta, state, 5); + if (test.flags & GdnObserved) { + ggml_set_output(gdn); + } + const size_t attention_bytes = channels * test.tokens * test.sequences * sizeof(float); + ggml_tensor * attention = + ggml_view_4d(ctx, gdn, width, heads, test.tokens, test.sequences, width * sizeof(float), + channels * sizeof(float), channels * test.tokens * sizeof(float), 0); + if (test.flags & AttentionObserved) { + ggml_set_output(attention); + } + ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, test.sequences, written, state_bytes, + state_bytes * test.sequences, attention_bytes); + ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, test.sequences, written, state_bytes, + 4 * state_bytes, (test.flags & Unaligned) ? sizeof(float) : state_bytes); + ggml_tensor * norm_input = (test.flags & AttentionReshape) ? + ggml_reshape_2d(ctx, attention, width, heads * test.tokens * test.sequences) : + attention; + ggml_tensor * rms_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); + ggml_tensor * raw_gate = + (test.flags & CachedGate) ? + ggml_view_4d(ctx, cache, width, heads, test.tokens, test.sequences, width * sizeof(float), + channels * sizeof(float), channels * test.tokens * sizeof(float), 0) : + ggml_scale(ctx, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, heads, test.tokens, test.sequences), + 0.5f); + ggml_tensor * shaped_gate = (test.flags & AttentionReshape) ? + ggml_reshape_2d(ctx, raw_gate, width, heads * test.tokens * test.sequences) : + raw_gate; + ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, norm_input, 1.e-6f), rms_weight); + ggml_tensor * gated = + ggml_mul(ctx, norm, (test.flags & TanhGate) ? ggml_tanh(ctx, shaped_gate) : ggml_silu(ctx, shaped_gate)); + ggml_tensor * projection_input = ggml_reshape_2d(ctx, gated, channels, test.tokens * test.sequences); + ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, test.type, channels, 4096); + ggml_tensor * projection = ggml_mul_mat(ctx, projection_weight, projection_input); + ggml_cgraph * graph = ggml_new_graph(ctx); + if (!(test.flags & UnreadyInput)) { + ggml_build_forward_expand(graph, raw); + } + ggml_build_forward_expand(graph, alpha); + ggml_build_forward_expand(graph, beta_raw); + if (!(test.flags & UnreadyGate)) { + ggml_build_forward_expand(graph, raw_gate); + } + ggml_build_forward_expand(graph, gathered); + if (test.flags & CacheRead) { + ggml_build_forward_expand(graph, ggml_scale(ctx, cache, 0.25f)); + } + ggml_build_forward_expand(graph, ggml_cpy(ctx, snapshots, target)); + ggml_build_forward_expand(graph, ggml_scale(ctx, projection, 0.5f)); + if (test.flags & GatherSide) { + ggml_build_forward_expand(graph, ggml_scale(ctx, gathered, 0.5f)); + } + if (test.flags & AttentionSide) { + ggml_build_forward_expand(graph, ggml_scale(ctx, attention, 0.5f)); + } + if (test.flags & OutputSide) { + 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 fused = 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_f32_wmma_head128_selected_snapshot_projection_rms_gate_q8"; + }); + REQUIRE((fused != plan.dispatches.end()) == test.fused); + if (test.fused) { + REQUIRE(fused->bindings.size() == 14); + require_compile_parameter(*fused, "llm.gated_delta_net.state_row_count", "20"); + REQUIRE(fused->bindings[7].value == imported.graph.values().find_tensor(cache)->id); + REQUIRE(fused->bindings[10].value == imported.graph.values().find_tensor(ids)->id); + const auto * output = imported.graph.values().find_tensor(gated); + const size_t bytes = static_cast(test.tokens * heads) * ggml_row_size(GGML_TYPE_Q8_1, width); + const auto * alternate = plan.metadata.find_alternate_value(output->id, GGML_TYPE_Q8_1, bytes); + REQUIRE(alternate != nullptr && alternate->alternate_value == fused->bindings.back().value); + 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_publish"; + })); + } + const auto commands = + ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + REQUIRE(commands.valid()); + if (!(test.flags & UnreadyGate)) { + REQUIRE(command_program_verifies(commands)); + } + } + ggml_free(ctx); +} + +static void register_hrx_backend_host_cases(test_runner::Suite & suite) { + suite.host_case("status", [] { run_status_checks(); }); + suite.host_case("command_plan_metadata", [] { run_command_plan_metadata_checks(); }); + suite.host_case("dispatch_registry", [] { run_dispatch_registry_checks(); }); + suite.host_case("graph_import", [] { run_graph_import_checks(); }); + suite.host_case("graph_import_mixed_backend_boundary", [] { run_graph_import_mixed_backend_boundary_checks(); }); + suite.host_case("graph_view_external_use", [] { run_graph_view_external_use_checks(); }); + suite.host_case("scale_f32_dispatch", [] { run_scale_f32_dispatch_checks(); }); + suite.host_case("scale_add_f32_dispatch", [] { run_scale_add_f32_dispatch_checks(); }); + suite.host_case("rmsnorm_two_binary_dispatch", [] { run_rmsnorm_two_binary_dispatch_checks(); }); + suite.host_case("rmsnorm_lowtoken_striped_dispatch", [] { run_rmsnorm_lowtoken_striped_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(); }); + suite.host_case("graph_index", [] { run_graph_index_checks(); }); + suite.host_case("rope_set_rows_dispatch", [] { run_rope_set_rows_dispatch_checks(); }); + suite.host_case("graph_traversal", [] { run_graph_traversal_checks(); }); + suite.host_case("qwen_token_embedding_dispatch", [] { run_qwen_token_embedding_dispatch_checks(); }); + suite.host_case("get_rows_rmsnorm_binary_dispatch", [] { run_get_rows_rmsnorm_binary_dispatch_checks(); }); + suite.host_case("get_rows_scale_dispatch", [] { run_get_rows_scale_dispatch_checks(); }); + suite.host_case("gather_add_dispatch", [] { run_gather_add_dispatch_checks(); }); + suite.host_case("qwen_flash_attention_dispatch", [] { run_qwen_flash_attention_dispatch_checks(); }); + suite.host_case("tiled_pair_matmul_postops_dispatch", [] { run_tiled_pair_matmul_postops_dispatch_checks(); }); + suite.host_case("qwen_attention_postprocess_dispatch", [] { run_qwen_attention_postprocess_dispatch_checks(); }); + suite.host_case("qwen_matmul_dispatch", [] { run_qwen_matmul_dispatch_checks(); }); + suite.host_case("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(); }); + suite.host_case("rmsnorm_binary_f16_output", [] { run_rmsnorm_binary_f16_output_checks(); }); + suite.host_case("swiglu_q8_output", [] { run_swiglu_q8_output_checks(); }); + suite.host_case("rmsnorm_gate_packed_output", [] { run_rmsnorm_gate_packed_output_checks(); }); + suite.host_case("rmsnorm_gate_q8_output", [] { run_rmsnorm_gate_q8_output_checks(); }); + suite.host_case("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(); }); + suite.host_case("quantized_conv4_dispatch", [] { run_quantized_conv4_dispatch_checks(); }); + suite.host_case("lfm_dconv3_dispatch", [] { run_lfm_dconv3_dispatch_checks(); }); + suite.host_case("gdn_rmsnorm_gate_dispatch", [] { run_gdn_rmsnorm_gate_dispatch_checks(); }); + suite.host_case("gdn_native_projection_pair_dispatch", [] { run_gdn_native_projection_pair_dispatch_checks(); }); + suite.host_case("gdn_selected_snapshot_dispatch", [] { run_gdn_selected_snapshot_dispatch_checks(); }); + suite.host_case("gdn_selected_rms_q8_dispatch", [] { run_gdn_selected_rms_q8_dispatch_checks(); }); + suite.host_case("llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); }); + suite.host_case("quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); }); + suite.host_case("prefill_value_projection_cache_dispatch", + [] { run_prefill_value_projection_cache_dispatch_checks(); }); + suite.host_case("qwen_terminal_q6k_q8_command.tokens1", [] { schedule_qwen_terminal_q6k_q8_command(1); }); + suite.host_case("qwen_terminal_q6k_q8_command.tokens18", [] { schedule_qwen_terminal_q6k_q8_command(18); }); + suite.host_case("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", + [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); }); + suite.host_case("qwen_router_top8_dispatch", [] { run_qwen_router_top8_dispatch_checks(); }); + suite.host_case("common_mul_mat_id_swiglu_dispatch", [] { run_common_mul_mat_id_swiglu_dispatch_checks(); }); + suite.host_case("qwen_routed_gate_up_dispatch", [] { run_qwen_routed_gate_up_dispatch_checks(); }); + suite.host_case("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(); }); + suite.host_case("layout_alias_scheduler_elision", [] { run_layout_alias_scheduler_elision_checks(); }); + suite.host_case("zero_output_scheduler_elision", [] { run_zero_output_scheduler_elision_checks(); }); + suite.host_case("transient_import", [] { run_transient_import_checks(); }); + suite.host_case("graph_view_preserves_shared_output_storage", + [] { run_graph_view_preserves_shared_output_storage_checks(); }); + suite.host_case("chained_dispatch_requires_transients", [] { run_chained_dispatch_requires_transients(); }); + suite.host_case("graph_replay_host_staging_is_not_ineligible", + [] { run_graph_replay_host_staging_is_not_ineligible(); }); + suite.host_case("multiple_transient_plan", [] { run_multiple_transient_plan_checks(); }); + suite.host_case("disjoint_transient_plan_packing", [] { run_disjoint_transient_plan_packing_checks(); }); + suite.host_case("command_program_kernel_dump", [] { run_command_program_kernel_dump_checks(); }); + suite.host_case("command_shape_hash", [] { run_command_shape_hash_checks(); }); + suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); + suite.host_case("graph_match_hash_collision", [] { run_graph_match_hash_collision_checks(); }); + suite.host_case("graph_match_bijection", [] { run_graph_match_bijection_checks(); }); + suite.host_case("validated_graph_uid_match", [] { run_validated_graph_uid_match_checks(); }); + suite.host_case("graph_executor_contract", [] { run_graph_executor_contract_checks(); }); + suite.host_case("loom_async_jit_environment", [] { + REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); + }); +} + +static void register_hrx_backend_device_cases(test_runner::Suite & suite) { + suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); + suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); + suite.device_case("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(); }); + suite.device_case("rope_scale_f32_numerics", [] { run_rope_scale_f32_numerics(); }); + suite.device_case("gather_add_rmsnorm_f32_numerics", [] { run_gather_add_rmsnorm_f32_numerics(); }); + suite.device_case("scale_f32", [] { run_scale_f32(); }); + suite.device_case("scale_f32_inplace", [] { run_scale_f32_inplace(); }); + suite.device_case("two_independent_add_f32", [] { run_two_independent_add_f32(); }); + suite.device_case("chained_add_f32", [] { run_chained_add_f32(); }); + suite.device_case("same_uid_distinct_graph_reuses_graph_program.default", + [] { run_same_uid_distinct_graph_reuses_graph_program(); }); + suite.device_case("same_uid_distinct_graph_reuses_graph_program.distinct_uid", + [] { run_same_uid_distinct_graph_reuses_graph_program(true); }); + suite.device_case("unsupported_op_fails", [] { run_unsupported_op_fails(); }); +} + +static void register_hrx_backend_cases(test_runner::Suite & suite) { + register_hrx_backend_host_cases(suite); + register_hrx_backend_device_cases(suite); +} + +int main(int argc, char ** argv) { + test_runner::Suite suite( + test_runner::Config::with_prefix("HRX backend test", "hrx-backend", "GGML_HRX_BACKEND_TEST")); + register_hrx_backend_cases(suite); + + const bool has_device = ggml_backend_hrx_get_device_count() != 0; + return suite.run(argc, argv, has_device); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4098acaaf91a..d55e44f11b2a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3943,6 +3943,24 @@ struct test_ssm_conv : public test_case { } }; +// GGML_OP_SSM_CONV + GGML_OP_MUL +struct test_ssm_conv_mul : public test_ssm_conv { + using test_ssm_conv::test_ssm_conv; + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "SSM_CONV_MUL"; + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); + ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne_b.data()); + ggml_tensor * out = ggml_ssm_conv(ctx, a, b); + ggml_tensor * rhs = ggml_new_tensor_3d(ctx, type, out->ne[0], out->ne[1], out->ne[2]); + return ggml_mul(ctx, out, rhs); + } +}; + // GGML_OP_SSM_CONV + GGML_OP_ADD (channel-wise bias, optional) + GGML_OP_UNARY(SILU) (fused operation) struct test_ssm_conv_bias_silu : public test_case { const ggml_type type; @@ -8723,6 +8741,8 @@ static std::vector> make_test_cases_eval() { // in-place tests test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 5, 4, 3}, false, 1e-6f, true)); + test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 32, 2, 1}, false, 1e-6f)); + test_cases.emplace_back(new test_rms_norm(GGML_TYPE_F32, {64, 32, 22, 1}, false, 1e-6f)); for (float eps : { 0.0f, 1e-6f, 1e-4f, 1e-1f, 1.0f }) { for (uint32_t n : { 64, 1025 }) { @@ -8766,6 +8786,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv - 1 + 64, d_inner, 4, 1}, {d_conv, d_inner, 1, 1})); } } + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {4, 2048, 1, 1}, {3, 2048, 1, 1})); + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {24, 2048, 1, 1}, {3, 2048, 1, 1})); + test_cases.emplace_back(new test_ssm_conv_mul(GGML_TYPE_F32, {4, 2048, 1, 1}, {3, 2048, 1, 1})); + test_cases.emplace_back(new test_ssm_conv_mul(GGML_TYPE_F32, {24, 2048, 1, 1}, {3, 2048, 1, 1})); // fused ssm_conv + (optional) bias_add + silu. The bias-only graph (no silu) is intentionally // not tested since there's no fusion for that pattern in ggml_cuda_can_fuse. @@ -9315,6 +9339,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_concat(GGML_TYPE_I64, {11, 12, 13, 14}, 7, dim, v)); } } + for (int v : { 0, 1, 2, 3 }) { + test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {64, 40, 2, 1}, 32, 0, v)); + test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {2, 2048, 1, 1}, 22, 0, v)); + } for (ggml_type type_a : { GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0 }) { for (int v : { 0, 4, 8, 12 }) { diff --git a/tests/test-hrx-buffer.cpp b/tests/test-hrx-buffer.cpp index 6d7d25703b03..e281a3199f09 100644 --- a/tests/test-hrx-buffer.cpp +++ b/tests/test-hrx-buffer.cpp @@ -6,6 +6,7 @@ #include "ggml.h" #include "hrx-interop-utils.h" #include "runtime/host-memory.h" +#include "testing_suite.h" #include #include @@ -169,10 +170,8 @@ static void run_host_buffer_checks(ggml_backend_t backend) { tensor->data = ggml_backend_buffer_get_base(local); REQUIRE(ggml_backend_buffer_init_tensor(local, tensor) == GGML_STATUS_SUCCESS); - const uint64_t upload_fallbacks = - context->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed); - const uint64_t download_fallbacks = - context->device->synchronous_download_fallbacks.load(std::memory_order_relaxed); + const uint64_t upload_fallbacks = context->device->synchronous_upload_fallbacks.load(std::memory_order_relaxed); + const uint64_t download_fallbacks = context->device->synchronous_download_fallbacks.load(std::memory_order_relaxed); auto * host_words = static_cast(ggml_backend_buffer_get_base(buffer)); for (size_t i = 0; i < 64; ++i) { host_words[i] = static_cast(i * 13 + 7); @@ -230,8 +229,8 @@ static void run_host_transfer_checks(ggml_backend_hrx_context * context) { hrx_synchronous_h2d(context->device->device, device_values.data(), staging.buffer, 0, device_values.size())); std::array download_result = {}; - REQUIRE(transfers.download_synchronous( - context->stream, staging.buffer, 12, download_result.data() + 4, 20).success()); + REQUIRE( + transfers.download_synchronous(context->stream, staging.buffer, 12, download_result.data() + 4, 20).success()); stats = transfers.stats(); REQUIRE(stats.downloads == 1); REQUIRE(stats.download_bytes == 20); @@ -339,7 +338,6 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { ggml::hrx::HostWeightCacheStats weight_stats = weights.stats(); REQUIRE(weight_stats.hits == 1); REQUIRE(weight_stats.misses == 3); - REQUIRE(weight_stats.layout_conflicts == 1); REQUIRE(weight_stats.allocation_count == 3); REQUIRE(weight_stats.resident_bytes == 96); @@ -351,27 +349,52 @@ static void run_host_weight_cache_checks(ggml_backend_hrx_context * context) { weight_stats = weights.stats(); REQUIRE(weight_stats.hits == 0); REQUIRE(weight_stats.misses == 0); - REQUIRE(weight_stats.layout_conflicts == 0); REQUIRE(weight_stats.allocation_count == 0); REQUIRE(weight_stats.resident_bytes == 0); } -int main() { - if (ggml_backend_hrx_get_device_count() == 0) { - std::fprintf(stderr, "test skipped: no HRX devices available\n"); - return 0; - } - +template +static void run_with_hrx_backend(F && f) { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); ggml_backend_hrx_context * context = backend_context(backend); + f(backend, context); + ggml_backend_free(backend); +} - run_backend_buffer_checks(backend); - run_host_buffer_checks(backend); - run_host_transfer_checks(context); - run_host_staging_checks(context); - run_host_weight_cache_checks(context); +static void register_hrx_buffer_cases(test_runner::Suite & suite) { + suite.device_case("backend_buffer", [] { + run_with_hrx_backend([](ggml_backend_t backend, ggml_backend_hrx_context *) { + run_backend_buffer_checks(backend); + }); + }); + suite.device_case("host_buffer", [] { + run_with_hrx_backend([](ggml_backend_t backend, ggml_backend_hrx_context *) { + run_host_buffer_checks(backend); + }); + }); + suite.device_case("host_transfer", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_host_transfer_checks(context); + }); + }); + suite.device_case("host_staging", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_host_staging_checks(context); + }); + }); + suite.device_case("host_weight_cache", [] { + run_with_hrx_backend([](ggml_backend_t, ggml_backend_hrx_context * context) { + run_host_weight_cache_checks(context); + }); + }); +} - ggml_backend_free(backend); - return 0; +int main(int argc, char ** argv) { + test_runner::Suite suite(test_runner::Config::with_prefix( + "HRX buffer test", "hrx-buffer", "GGML_HRX_BUFFER_TEST", 1)); + register_hrx_buffer_cases(suite); + + const bool has_device = ggml_backend_hrx_get_device_count() != 0; + return suite.run(argc, argv, has_device); } diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index d8480a3492bf..d90ae6a4ee98 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -3,6 +3,7 @@ #include "kernel-corpus/kernel-corpus.h" #include "runtime/kernel-executable-cache.h" #include "runtime/loom-kernel-jit.h" +#include "testing_suite.h" #include #include @@ -130,10 +131,30 @@ static const ggml::hrx::KernelDefinition * find_targeted_kernel(const char * nam return nullptr; } -static ggml::hrx::Dispatch make_add_dispatch(const ggml::hrx::KernelDefinition & definition, int64_t element_count) { +static std::map binary_f32_exact_config(const char * op, int64_t element_count) { + const std::string count = std::to_string(element_count); + return { + { "ggml.binary_f32.op", op }, + { "ggml.binary_f32.ne0", count }, + { "ggml.binary_f32.ne1", "1" }, + { "ggml.binary_f32.ne2", "1" }, + { "ggml.binary_f32.src0_stride1", count }, + { "ggml.binary_f32.src0_stride2", count }, + { "ggml.binary_f32.src0_stride3", count }, + { "ggml.binary_f32.src1_stride1", count }, + { "ggml.binary_f32.src1_stride2", count }, + { "ggml.binary_f32.src1_stride3", count }, + { "ggml.binary_f32.src0_span", count }, + { "ggml.binary_f32.src1_span", count }, + }; +} + +static ggml::hrx::Dispatch make_binary_add_dispatch(const ggml::hrx::KernelDefinition & definition, + int64_t element_count) { ggml::hrx::Dispatch dispatch; dispatch.kernel.kernel_id = definition.id; dispatch.kernel.integer_parameters.emplace("element_count", element_count); + dispatch.kernel.compile_parameters = binary_f32_exact_config("0", element_count); dispatch.bindings.reserve(definition.bindings.size()); for (size_t i = 0; i < definition.bindings.size(); ++i) { dispatch.bindings.push_back({ @@ -145,6 +166,139 @@ static ggml::hrx::Dispatch make_add_dispatch(const ggml::hrx::KernelDefinition & return dispatch; } +static std::map binary_f32_exact_workload(int64_t element_count) { + return { + { "element_count", element_count }, + }; +} + +static std::map binary_bc_f32_workload() { + return { + { "element_count", 48 }, + { "ne0", 16 }, + { "ne1", 3 }, + { "ne2", 1 }, + { "ne3", 1 }, + { "src0_element_count", 48 }, + { "src1_element_count", 16 }, + }; +} + +static std::map binary_bc_f32_config(const char * op) { + return { + { "ggml.binary_bc_f32.op", op }, + { "ggml.binary_bc_f32.src0_broadcast_dim0", "0" }, + { "ggml.binary_bc_f32.src0_broadcast_dim1", "0" }, + { "ggml.binary_bc_f32.src0_broadcast_dim2", "0" }, + { "ggml.binary_bc_f32.src0_broadcast_dim3", "0" }, + { "ggml.binary_bc_f32.src1_broadcast_dim0", "0" }, + { "ggml.binary_bc_f32.src1_broadcast_dim1", "1" }, + { "ggml.binary_bc_f32.src1_broadcast_dim2", "0" }, + { "ggml.binary_bc_f32.src1_broadcast_dim3", "0" }, + }; +} + +static std::map token_count_workload(int64_t token_count) { + return { + { "token_count", token_count }, + }; +} + +static std::map moe_routing_config() { + return { + { "ggml.moe_routing.route_count", "8" }, + { "ggml.moe_routing.expert_count", "128" }, + { "ggml.moe_routing.descriptor_expert_mask", "127" }, + { "ggml.moe_routing.descriptor_partition_shift", "7" }, + { "ggml.moe_routing.descriptor_row_count_shift", "13" }, + { "ggml.moe_routing.partition_workgroup_size", "128" }, + { "ggml.workload.token_capacity", "1" }, + }; +} + +static std::map flash_attention_workload(int64_t query_token_count, + int64_t key_value_token_count) { + return { + { "query_token_count", query_token_count }, + { "key_value_token_count", key_value_token_count }, + }; +} + +static std::map flash_attention_decode_workload(int64_t key_value_token_count) { + return { + { "key_value_token_count", key_value_token_count }, + }; +} + +static std::map flash_attention_config(const char * query_head_count, + const char * key_value_head_count, + const char * head_size, + const char * attention_scale) { + return { + { "ggml.flash_attention.query_head_count", query_head_count }, + { "ggml.flash_attention.key_value_head_count", key_value_head_count }, + { "ggml.flash_attention.qk_head_size", head_size }, + { "ggml.flash_attention.value_head_size", head_size }, + { "ggml.flash_attention.attention_scale", attention_scale }, + { "ggml.flash_attention.apply_gate", "0" }, + { "ggml.flash_attention.gate_stride_head", "1" }, + { "ggml.flash_attention.gate_stride_token", "1" }, + }; +} + +static std::map flash_attention_decode_base_config(const char * query_head_count, + const char * key_value_head_count, + const char * head_size, + const char * attention_scale) { + return { + { "ggml.flash_attention.query_head_count", query_head_count }, + { "ggml.flash_attention.key_value_head_count", key_value_head_count }, + { "ggml.flash_attention.qk_head_size", head_size }, + { "ggml.flash_attention.value_head_size", head_size }, + { "ggml.flash_attention.attention_scale", attention_scale }, + { "ggml.flash_attention.apply_gate", "0" }, + { "ggml.flash_attention.gate_stride_head", head_size }, + { "ggml.flash_attention.gate_stride_token", head_size }, + }; +} + +static std::map flash_attention_decode_config(const char * query_head_count, + const char * key_value_head_count, + const char * head_size, + const char * attention_scale, + const char * key_value_token_capacity) { + std::map config = + flash_attention_decode_base_config(query_head_count, key_value_head_count, head_size, attention_scale); + config["ggml.flash_attention.decode.key_value_token_capacity"] = key_value_token_capacity; + return config; +} + +static std::map mul_mat_id_f16_f16_config(const char * weight_format) { + return { + { "ggml.mul_mat_id_f16_f16.input_size", "768" }, + { "ggml.mul_mat_id_f16_f16.route_count", "8" }, + { "ggml.mul_mat_id_f16_f16.expert_count", "128" }, + { "ggml.mul_mat_id_f16_f16.output_size", "2048" }, + { "ggml.mul_mat_id_f16_f16.weight_format", weight_format }, + { "ggml.workload.token_capacity", "4" }, + }; +} + +static std::map mul_mat_id_swiglu_f16_f16_config(const char * gate_format, + const char * up_format) { + return { + { "ggml.mul_mat_id_swiglu_f16_f16.input_size", "2048" }, + { "ggml.mul_mat_id_swiglu_f16_f16.output_size", "768" }, + { "ggml.mul_mat_id_swiglu_f16_f16.expert_count", "128" }, + { "ggml.mul_mat_id_swiglu_f16_f16.route_count", "8" }, + { "ggml.mul_mat_id_swiglu_f16_f16.gate_weight_format", gate_format }, + { "ggml.mul_mat_id_swiglu_f16_f16.up_weight_format", up_format }, + { "ggml.mul_mat_id_swiglu_f16_f16.descriptor_expert_mask", "127" }, + { "ggml.mul_mat_id_swiglu_f16_f16.descriptor_partition_shift", "7" }, + { "ggml.mul_mat_id_swiglu_f16_f16.descriptor_row_count_shift", "13" }, + }; +} + static ggml::hrx::LoomKernelCompileRequest make_compile_request( const ggml::hrx::KernelDefinition & definition, const std::map & workload, @@ -240,36 +394,131 @@ static int64_t elapsed_us(std::chrono::steady_clock::time_point begin, std::chro static void run_cache_materialize_case(HrxTestDevice & device, ggml::hrx::LoomJitMode mode, - const ggml::hrx::KernelDefinition & add) { + const ggml::hrx::KernelDefinition & binary) { ggml::hrx::KernelExecutablePrepareContext context = {}; context.device = device.device; context.target = device.architecture.c_str(); - ggml::hrx::Dispatch dispatch = make_add_dispatch(add, 64); + 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, add, 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, add, 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"; - std::printf("%s KernelExecutableCache materialized ggml_add_f32 for %s\n", mode_name, device.architecture.c_str()); + std::printf("%s KernelExecutableCache materialized ggml_binary_f32 for %s\n", mode_name, + 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) { @@ -294,25 +543,34 @@ 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()); } } // namespace -int main() { +static void run_loom_jit_compile_checks() { static constexpr const char * kTarget = "gfx1100"; - const ggml::hrx::KernelDefinition & add = find_kernel("ggml_add_f32"); + 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("qwen3_moe_build_expert_table"); - const ggml::hrx::KernelDefinition & partition_table = find_kernel("qwen3_moe_build_expert_partition_table"); + 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"); + const ggml::hrx::KernelDefinition & mul_mat_id_f16 = find_kernel("ggml_mul_mat_id_f16_f16_wmma"); + const ggml::hrx::KernelDefinition & swiglu_f16 = find_kernel("ggml_mul_mat_id_swiglu_f16_f16_wmma"); + const ggml::hrx::KernelDefinition & flash_prefill = find_kernel("ggml_flash_attention_f32_f16_wmma"); + const ggml::hrx::KernelDefinition & decode_mul_mat = find_kernel("ggml_mul_mat_f32_f32_decode_wave64"); + const ggml::hrx::KernelDefinition & flash_decode = + find_kernel("ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); std::string sync_error; std::unique_ptr sync_jit = @@ -321,10 +579,8 @@ int main() { REQUIRE(!sync_jit->async_enabled()); const auto sync_begin = std::chrono::steady_clock::now(); - ggml::hrx::LoomCompiledKernelRef sync_ref = compile_kernel(*sync_jit, "sync-add-64", add, - { - { "element_count", 64 } - }); + ggml::hrx::LoomCompiledKernelRef sync_ref = compile_kernel( + *sync_jit, "sync-binary-add-64", binary, binary_f32_exact_workload(64), binary_f32_exact_config("0", 64)); require_compiled_kernel(sync_ref); const auto sync_end = std::chrono::steady_clock::now(); const int64_t sync_compile_us = elapsed_us(sync_begin, sync_end); @@ -337,33 +593,34 @@ int main() { REQUIRE(async_jit->async_enabled()); std::vector refs; - refs.reserve(9); + refs.reserve(31); const auto enqueue_begin = std::chrono::steady_clock::now(); - refs.push_back(compile_kernel(*async_jit, "async-add-64", add, - { - { "element_count", 64 } - })); - refs.push_back(compile_kernel(*async_jit, "async-add-128", add, - { - { "element_count", 128 } - })); - refs.push_back(compile_kernel(*async_jit, "async-add-256", add, - { - { "element_count", 256 } - })); - refs.push_back(compile_kernel(*async_jit, "async-add-512", add, - { - { "element_count", 512 } - })); + refs.push_back(compile_kernel(*async_jit, "async-binary-add-64", binary, binary_f32_exact_workload(64), + binary_f32_exact_config("0", 64))); + refs.push_back(compile_kernel(*async_jit, "async-binary-add-128", binary, binary_f32_exact_workload(128), + binary_f32_exact_config("0", 128))); + refs.push_back(compile_kernel(*async_jit, "async-binary-add-256", binary, binary_f32_exact_workload(256), + binary_f32_exact_config("0", 256))); + refs.push_back(compile_kernel(*async_jit, "async-binary-add-512", binary, binary_f32_exact_workload(512), + binary_f32_exact_config("0", 512))); + refs.push_back(compile_kernel(*async_jit, "async-binary-geglu-256", binary, binary_f32_exact_workload(256), + binary_f32_exact_config("5", 256))); + refs.push_back(compile_kernel(*async_jit, "async-binary-reglu-256", binary, binary_f32_exact_workload(256), + binary_f32_exact_config("6", 256))); + refs.push_back(compile_kernel(*async_jit, "async-binary-geglu-erf-256", binary, binary_f32_exact_workload(256), + binary_f32_exact_config("7", 256))); + refs.push_back(compile_kernel(*async_jit, "async-binary-geglu-quick-256", binary, binary_f32_exact_workload(256), + binary_f32_exact_config("8", 256))); + refs.push_back(compile_kernel(*async_jit, "async-binary-bc-mul-48", binary_bc, binary_bc_f32_workload(), + binary_bc_f32_config("2"))); refs.push_back(compile_kernel(*async_jit, "async-rmsnorm-1", rmsnorm, { { "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, { @@ -382,29 +639,64 @@ int main() { { "route_stride", 8 }, { "expert_count", 128 }, }, - { - { "qwen3_moe.routed_gate_up.expert_count", "128" }, - { "qwen3_moe.routed_gate_up.route_count", "8" }, - { "qwen3_moe.workload.token_capacity", "1" }, - })); + moe_routing_config())); refs.push_back(compile_kernel(*async_jit, "async-partition-table-1", partition_table, { { "token_count", 1 }, { "route_count", 8 }, { "expert_count", 128 }, }, - { - { "qwen3_moe.routed_gate_up.expert_count", "128" }, - { "qwen3_moe.routed_gate_up.route_count", "8" }, - { "qwen3_moe.workload.token_capacity", "1" }, - })); + moe_routing_config())); // Exercise gather-add coverage in the async JIT path. refs.push_back(compile_kernel(*async_jit, "async-gather-add-2-to-1", gather_add, { { "source_token_count", 2 }, { "output_token_count", 1 }, { "hidden_size", 2048 }, + }, + { + { "ggml.gather_add_f32.binary_op", "3" }, })); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-f16-q4", mul_mat_id_f16, token_count_workload(4), + mul_mat_id_f16_f16_config("4"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-f16-q6", mul_mat_id_f16, token_count_workload(4), + mul_mat_id_f16_f16_config("6"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-q4k", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("4", "4"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-q6k", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("6", "6"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-iq4-xs", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("23", "23"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-q3k", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("11", "11"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-iq4-nl", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("20", "20"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-q8-0", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("80", "80"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-q8-1", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("81", "81"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-swiglu-f16-f16", swiglu_f16, token_count_workload(4), + mul_mat_id_swiglu_f16_f16_config("16", "16"))); + refs.push_back(compile_kernel(*async_jit, "async-mul-mat-decode-q8-0-gemma-vocab", decode_mul_mat, + { + { "token_count", 1 }, + { "input_size", 3840 }, + { "output_size", 262208 }, + }, + { + { "ggml.mul_mat_f32_f32_decode.token_capacity", "1" }, + { "ggml.mul_mat_f32_f32_decode.output_capacity", "262208" }, + { "ggml.mul_mat_f32_f32_decode.weight_format", "80" }, + })); + refs.push_back(compile_kernel(*async_jit, "async-flash-prefill-128", flash_prefill, + flash_attention_workload(128, 128), + flash_attention_config("32", "4", "128", "0.0883883461"))); + refs.push_back(compile_kernel(*async_jit, "async-flash-prefill-256", flash_prefill, + flash_attention_workload(128, 128), + flash_attention_config("32", "4", "256", "0.0625"))); + refs.push_back(compile_kernel(*async_jit, "async-flash-decode-128", flash_decode, + flash_attention_decode_workload(64), + flash_attention_decode_config("32", "4", "128", "0.0883883461", "64"))); const auto enqueue_end = std::chrono::steady_clock::now(); const int64_t enqueue_us = elapsed_us(enqueue_begin, enqueue_end); std::printf("async Loom enqueue completed in %ld us\n", static_cast(enqueue_us)); @@ -418,14 +710,861 @@ int main() { std::printf("async Loom JIT compiled %zu kernels\n", refs.size()); - HrxTestDevice device; - if (device.open()) { - run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, add); - run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, add); - run_targeted_export_materialize_case(device); - } else { - std::printf("skipping KernelExecutableCache materialization checks: no HRX device available\n"); + const auto & q6_prefill = find_kernel("ggml_mul_mat_q6_k_f16_wmma_prefill_wave32"); + + const struct { + int64_t tokens; + int64_t inputs; + int64_t outputs; + int64_t accumulate; + int64_t workgroup_size; + int64_t token_tiles; + int64_t output_tiles; + } q6_launch_cases[] = { + { 128, 256, 4096, 0, 256, 1, 32 }, + { 256, 256, 8192, 0, 512, 1, 64 }, + { 512, 256, 1024, 0, 512, 2, 8 }, + { 512, 256, 3840, 0, 512, 2, 30 }, + { 512, 256, 4096, 0, 1024, 1, 32 }, + { 512, 256, 4224, 0, 1024, 1, 33 }, + { 512, 256, 4096, 1, 256, 4, 32 }, + {1024, 256, 2048, 0, 1024, 2, 16 }, + { 512, 17408, 5120, 0, 512, 1, 80 }, + { 512, 5120, 10240, 0, 1024, 1, 80 }, + { 512, 5120, 1024, 0, 512, 2, 8 }, + { 512, 17408, 4096, 0, 512, 1, 64 }, + { 512, 5120, 1984, 0, 256, 4, 16 }, + { 512, 5120, 2048, 0, 512, 1, 32 }, + {1024, 5120, 1024, 0, 512, 2, 16 }, + {2048, 5120, 1024, 0, 512, 4, 16 }, + }; + + for (const auto & test : q6_launch_cases) { + const std::string key = "q6-prefill-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + "-" + + std::to_string(test.outputs) + "-" + std::to_string(test.accumulate); + auto ref = + compile_kernel(*sync_jit, key, q6_prefill, token_count_workload(test.tokens), + { + { "ggml.mul_mat_q6_k_packed.input_size", std::to_string(test.inputs) }, + { "ggml.mul_mat_q6_k_packed.output_size", std::to_string(test.outputs) }, + { "ggml.mul_mat_q6_k_packed.output_accumulation", std::to_string(test.accumulate) }, + { "ggml.mul_mat_q6_k_packed.weight_offset", "0" }, + { "ggml.mul_mat_q6_k_packed.token_capacity", std::to_string(test.tokens) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == test.workgroup_size); + REQUIRE(compiled.launch_config.workgroup_count[0] == test.token_tiles); + REQUIRE(compiled.launch_config.workgroup_count[1] == test.output_tiles); + compiled.reset(); + } + + const auto & q4_prefill = find_kernel("ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); + + const struct { + int64_t tokens; + int64_t inputs; + int64_t outputs; + int64_t workgroup_size; + int64_t token_tiles; + int64_t output_tiles; + int64_t input_layout; + } q4_launch_cases[] = { + { 256, 8192, 8192, 512, 1, 64, 0 }, + { 512, 8192, 1024, 512, 2, 8, 0 }, + { 512, 8192, 3840, 512, 2, 30, 0 }, + { 512, 8192, 4096, 512, 1, 64, 1 }, + { 512, 8192, 4224, 512, 1, 66, 1 }, + { 512, 8192, 12288, 512, 1, 192, 1 }, + { 512, 8192, 16384, 512, 1, 256, 1 }, + {1024, 8192, 2048, 1024, 4, 8, 0 }, + { 512, 17408, 5120, 1024, 2, 20, 0 }, + { 512, 5120, 6144, 512, 1, 96, 1 }, + { 512, 6144, 5120, 512, 1, 80, 1 }, + { 512, 5120, 10240, 512, 1, 160, 1 }, + { 512, 5120, 8192, 512, 1, 128, 1 }, + {1024, 5120, 4096, 1024, 4, 16, 0 }, + { 512, 17408, 5120, 512, 1, 80, 1 }, + }; + + for (const auto & test : q4_launch_cases) { + const std::string key = "q4-prefill-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + "-" + + std::to_string(test.outputs) + "-" + std::to_string(test.input_layout); + auto ref = compile_kernel(*sync_jit, key, q4_prefill, token_count_workload(test.tokens), + { + { "ggml.mul_mat.input_size", std::to_string(test.inputs) }, + { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, + { "ggml.workload.token_capacity", std::to_string(test.tokens) }, + { "ggml.mul_mat.f16_input_layout", std::to_string(test.input_layout) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == test.workgroup_size); + REQUIRE(compiled.launch_config.workgroup_count[0] == test.token_tiles); + REQUIRE(compiled.launch_config.workgroup_count[1] == test.output_tiles); + compiled.reset(); + } + + const auto & packed_f16 = find_kernel("ggml_copy_f16_k16_major"); + for (const auto * input_is_f16 : { "0", "1" }) { + auto ref = + compile_kernel(*sync_jit, std::string("packed-f16-") + input_is_f16, packed_f16, token_count_workload(512), + { + { "ggml.copy_f16_k16_major.input_size", "256" }, + { "ggml.copy_f16_k16_major.token_count", "512" }, + { "ggml.copy_f16_k16_major.input_is_f16", input_is_f16 }, + }); + require_compiled_kernel(ref); + } + + 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 }, + }, + { + }); + require_compiled_kernel(transpose_ref); + auto transposed_attention_config = flash_attention_config("24", "4", "256", "0.0625"); + transposed_attention_config.emplace("ggml.flash_attention.value_layout", "1"); + auto transposed_attention_ref = + compile_kernel(*sync_jit, "flash-transposed-value", flash_prefill, + { + { "query_token_count", 512 }, + { "key_value_token_count", 8192 } + }, + transposed_attention_config); + require_compiled_kernel(transposed_attention_ref); + + const auto & paired_q4 = find_kernel("ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"); + + const struct { + int64_t tokens; + int64_t outputs; + int64_t threads; + int64_t token_tiles; + } paired_q4_cases[] = { + { 256, 4096, 512, 1 }, + { 512, 3840, 512, 2 }, + { 512, 4096, 512, 1 }, + { 768, 8192, 512, 3 }, + {1024, 2048, 512, 2 }, + {2048,17408, 512, 4 }, + }; + + for (const auto & test : paired_q4_cases) { + const std::string key = "paired-q4-" + std::to_string(test.tokens) + "-" + std::to_string(test.outputs); + auto ref = compile_kernel(*sync_jit, key, paired_q4, token_count_workload(test.tokens), + { + { "ggml.mul_mat_swiglu.input_size", "5120" }, + { "ggml.mul_mat_swiglu.output_size", std::to_string(test.outputs) }, + { "ggml.mul_mat_swiglu.op", "4" }, + { "ggml.mul_mat_swiglu.f16_output_layout", "0" }, + { "ggml.workload.token_capacity", std::to_string(test.tokens) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == test.threads); + REQUIRE(compiled.launch_config.workgroup_count[0] == test.token_tiles); + const bool packed = test.tokens % 512 == 0 && (test.tokens / 512) * (test.outputs / 64) >= 64; + REQUIRE(compiled.launch_config.workgroup_count[1] == test.outputs / (packed ? 32 : 64)); + compiled.reset(); + } + + for (const int64_t tokens : { 256, 512, 1024, 2048 }) { + auto ref = compile_kernel(*sync_jit, "paired-q4-packed-output-" + std::to_string(tokens), paired_q4, + token_count_workload(tokens), + { + { "ggml.mul_mat_swiglu.input_size", "5120" }, + { "ggml.mul_mat_swiglu.output_size", "17408" }, + { "ggml.mul_mat_swiglu.op", "4" }, + { "ggml.mul_mat_swiglu.f16_output_layout", "1" }, + { "ggml.workload.token_capacity", std::to_string(tokens) }, + }); + require_compiled_kernel(ref); + } + + const auto & paired_q4_decode = find_kernel("ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot"); + + const struct { + int64_t tokens; + int64_t inputs; + int64_t outputs; + int64_t threads; + int64_t row_tiles; + } paired_q4_decode_cases[] = { + { 1, 5120, 17408, 128, 2176 }, + { 1, 4096, 8192, 128, 1024 }, + { 1, 4096, 16384, 128, 2048 }, + { 1, 8192, 8192, 128, 1024 }, + { 1, 5376, 17408, 128, 2176 }, + { 2, 5120, 17408, 512, 1088 }, + { 3, 5120, 17408, 512, 1088 }, + { 4, 5120, 17408, 512, 1088 }, + { 5, 5120, 17408, 512, 1088 }, + { 5, 4096, 8192, 512, 256 }, + { 5, 4096, 16384, 512, 1024 }, + { 3, 8192, 8192, 512, 512 }, + { 5, 5632, 12288, 512, 768 }, + { 5, 5376, 17408, 512, 544 }, + { 5, 6144, 17408, 128, 2176 }, + { 5, 256, 4032, 128, 504 }, + { 5, 256, 4096, 512, 128 }, + { 5, 5632, 4096, 512, 128 }, + { 5, 5888, 4096, 128, 512 }, + }; + + for (const auto & test : paired_q4_decode_cases) { + const std::string key = "paired-q4-decode-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + + "-" + std::to_string(test.outputs); + auto ref = compile_kernel(*sync_jit, key, paired_q4_decode, token_count_workload(test.tokens), + { + { "ggml.mul_mat_swiglu.input_size", std::to_string(test.inputs) }, + { "ggml.mul_mat_swiglu.output_size", std::to_string(test.outputs) }, + { "ggml.mul_mat_swiglu.op", "4" }, + { "ggml.workload.token_capacity", std::to_string(test.tokens) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == test.threads); + REQUIRE(compiled.launch_config.workgroup_count[0] == test.row_tiles); + REQUIRE(compiled.launch_config.workgroup_count[1] == 1); + compiled.reset(); + } + + const auto & quant_decode = find_kernel("ggml_mul_mat_f32_f32_wmma"); + + const struct { + int64_t tokens; + int64_t inputs; + int64_t outputs; + int64_t format; + int64_t threads; + int64_t row_tiles; + } quant_decode_cases[] = { + { 1, 5120, 48, 4, 256, 12 }, + { 2, 5120, 48, 4, 1024, 12 }, + { 3, 5120, 48, 4, 1024, 12 }, + { 4, 5120, 48, 4, 1024, 12 }, + { 5, 5120, 48, 4, 1024, 12 }, + { 5, 5120, 48, 5, 1024, 12 }, + { 5, 5120, 48, 6, 1024, 12 }, + { 5, 4096, 1, 4, 1024, 1 }, + { 5, 4096, 16, 4, 1024, 4 }, + { 5, 8192, 63, 4, 1024, 16 }, + { 5, 32768, 48, 4, 1024, 12 }, + { 5, 3840, 48, 4, 256, 12 }, + { 5, 4352, 48, 4, 256, 12 }, + { 5, 4096, 64, 4, 128, 8 }, + { 6, 5120, 48, 4, 128, 1 }, + { 1, 17408, 5120, 46, 128, 640 }, + { 4, 17408, 5120, 46, 128, 640 }, + { 5, 17408, 5120, 46, 512, 160 }, + { 5, 8192, 4096, 46, 512, 128 }, + { 5, 8448, 4160, 46, 512, 130 }, + { 5, 7936, 4096, 46, 128, 512 }, + { 5, 8192, 4032, 46, 128, 504 }, + { 5, 8192, 8256, 46, 128, 1032 }, + { 1, 17408, 5120, 44, 128, 640 }, + { 4, 17408, 5120, 44, 128, 640 }, + { 5, 17408, 5120, 44, 1024, 80 }, + { 5, 8192, 4096, 44, 128, 1024 }, + { 1, 6144, 5120, 44, 128, 640 }, + { 4, 6144, 5120, 44, 128, 640 }, + { 5, 6144, 5120, 44, 128, 1280 }, + { 5, 6400, 5120, 44, 128, 640 }, + { 5, 6656, 5120, 44, 128, 1280 }, + { 5, 6144, 4096, 44, 128, 512 }, + { 5, 6144, 6144, 44, 128, 768 }, + { 5, 5120, 6144, 44, 128, 768 }, + { 5, 8448, 4160, 44, 128, 520 }, + { 5, 9216, 9216, 44, 128, 1152 }, + { 5, 12288, 8192, 44, 128, 1024 }, + { 5, 16384, 4096, 44, 128, 512 }, + { 5, 16384, 5056, 44, 128, 632 }, + { 5, 16384, 5120, 44, 1024, 80 }, + { 5, 20480, 4096, 44, 1024, 64 }, + { 5, 20736, 4160, 44, 1024, 65 }, + { 5, 7936, 4096, 44, 128, 512 }, + { 5, 8192, 4032, 44, 128, 504 }, + { 5, 8192, 8256, 44, 128, 1032 }, + }; + + for (const auto & test : quant_decode_cases) { + const std::string key = "quant-decode-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + + "-" + std::to_string(test.outputs) + "-" + std::to_string(test.format); + auto ref = + compile_kernel(*sync_jit, key, quant_decode, token_count_workload(test.tokens), + { + { "ggml.mul_mat.input_size", std::to_string(test.inputs) }, + { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, + { "ggml.mul_mat.weight_format", std::to_string(test.format) }, + { "ggml.mul_mat.activation_format", test.format == 44 || test.format == 46 ? "9" : "0" }, + { "ggml.mul_mat.output_accumulation", "0" }, + { "ggml.mul_mat.output_unary_op", "23" }, + { "ggml.workload.token_capacity", std::to_string(test.tokens) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == test.threads); + REQUIRE(compiled.launch_config.workgroup_count[0] == test.row_tiles); + REQUIRE(compiled.launch_config.workgroup_count[1] == 1); + compiled.reset(); + } + + + for (const char * name : { "ggml_mul_mat_quantized_f16_wmma_prefill_conv4", + "ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" }) { + const auto & quantized_conv4 = find_kernel(name); + REQUIRE(quantized_conv4.bindings.size() == (std::string(name).find("interior") != std::string::npos ? 5 : 6)); + for (const int format : {4, 6}) { + for (const int64_t channels : {2112, 8192, 10240}) { + auto ref = compile_kernel(*sync_jit, std::string(name) + "-" + std::to_string(format) + "-" + + std::to_string(channels), quantized_conv4, token_count_workload(512), { + {"ggml.mul_mat.input_size", "5120"}, + {"ggml.mul_mat.output_size", std::to_string(channels)}, + {"ggml.mul_mat.weight_format", std::to_string(format)}, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == 512); + REQUIRE(compiled.launch_config.workgroup_count[0] == 1); + REQUIRE(compiled.launch_config.workgroup_count[1] == channels / 64); + compiled.reset(); + } + } } - return 0; + const auto & conv_finish = find_kernel("llm_ssm_conv_dconv4_silu_prefill_finish_f32"); + REQUIRE(conv_finish.bindings.size() == 5); + for (const int64_t channels : { 2112, 8192, 8256, 10240 }) { + auto ref = compile_kernel(*sync_jit, "conv-finish-" + std::to_string(channels), conv_finish, + { + }, + { + { "llm.ssm_conv.generic.d_inner", std::to_string(channels) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == 256); + REQUIRE(compiled.launch_config.workgroup_count[0] == (channels + 255) / 256); + } + + const auto & ssm_prefill = find_kernel("llm_ssm_conv_dconv4_silu_prefill_512_wg1024"); + for (const int64_t channels : { 8192, 8224, 8256, 10208, 10240 }) { + auto ref = compile_kernel(*sync_jit, "ssm-prefill-" + std::to_string(channels), ssm_prefill, + { + }, + { + { "llm.ssm_conv.prefill.d_conv", "4" }, + { "llm.ssm_conv.prefill.d_inner", std::to_string(channels) }, + { "llm.ssm_conv.prefill.n_t", "512" }, + { "llm.ssm_conv.prefill.n_s", "1" }, + { "llm.ssm_conv.prefill.state_row_stride", std::to_string(channels) }, + { "llm.ssm_conv.prefill.x_row_stride", std::to_string(channels) }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == 1024); + REQUIRE(compiled.launch_config.workgroup_count[0] == channels / (channels % 64 == 0 ? 64 : 32)); + REQUIRE(compiled.launch_config.workgroup_count[1] == 1); + compiled.reset(); + } + + const struct { + const char * name; + int64_t tokens; + size_t bindings; + } gdn_cases[] = { + { "llm_gated_delta_net_f32_wmma_head128", 17, 7 }, + { "llm_gated_delta_net_f32_wmma_head128_rmsnorm_gate", 17, 11 }, + { "llm_gated_delta_net_f32_wmma_head128_projection_epilogue", 16, 9 }, + { "llm_gated_delta_net_f32_wmma_head128_inplace", 1, 7 }, + { "llm_gated_delta_net_f32_wmma_head128_inplace_projection_epilogue", 1, 9 }, + { "llm_gated_delta_net_f32_wmma_head128_snapshot", 5, 8 }, + { "llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue", 5, 10 }, + }; + + for (const auto & test : gdn_cases) { + const auto & definition = find_kernel(test.name); + REQUIRE(definition.bindings.size() == test.bindings); + auto ref = compile_kernel(*sync_jit, test.name, definition, + { + }, + { + { "llm.gated_delta_net.head_width", "128" }, + { "llm.gated_delta_net.head_count", "3" }, + { "llm.gated_delta_net.token_count", std::to_string(test.tokens) }, + { "llm.gated_delta_net.sequence_count", "1" }, + { "llm.gated_delta_net.qk_stride1", "128" }, + { "llm.gated_delta_net.qk_stride2", "640" }, + { "llm.gated_delta_net.qk_stride3", std::to_string(test.tokens * 640) }, + { "llm.gated_delta_net.value_stride1", "128" }, + { "llm.gated_delta_net.value_stride2", "640" }, + { "llm.gated_delta_net.value_stride3", std::to_string(test.tokens * 640) }, + { "llm.gated_delta_net.scalar_stride1", "1" }, + { "llm.gated_delta_net.scalar_stride2", "3" }, + { "llm.gated_delta_net.scalar_stride3", std::to_string(test.tokens * 3) }, + { "llm.gated_delta_net.query_head_count", "1" }, + { "llm.gated_delta_net.query_sequence_ratio", "1" }, + { "llm.gated_delta_net.snapshot_stride", "49216" }, + { "llm.gated_delta_net.l2_epsilon", "0.000001" }, + { "ggml.rmsnorm_gate_f32.rms_epsilon", "0.000001" }, + { "ggml.rmsnorm_gate_f32.gate_op", "15" }, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == 256); + REQUIRE(compiled.launch_config.workgroup_count[0] == 3); + REQUIRE(compiled.launch_config.workgroup_count[1] == 1); + compiled.reset(); + } +} + +static void run_moe_routing_compile_checks() { + static constexpr const char * kTarget = "gfx1100"; + + std::string error; + std::unique_ptr jit = ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); + REQUIRE(jit != nullptr); + + const std::map expert_table_workload = { + { "token_count", 1 }, + { "route_count", 8 }, + { "route_stride", 8 }, + { "expert_count", 128 }, + }; + const auto & expert_table = find_kernel("ggml_moe_build_expert_table"); + require_compiled_kernel( + compile_kernel(*jit, "moe-expert-table", expert_table, expert_table_workload, moe_routing_config())); + + const std::map partition_table_workload = { + { "token_count", 1 }, + { "route_count", 8 }, + { "expert_count", 128 }, + }; + 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() { + static constexpr const char * kTarget = "gfx1100"; + + std::string error; + std::unique_ptr jit = ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); + REQUIRE(jit != nullptr); + + const auto & rmsnorm_f32_f16 = find_kernel("ggml_rmsnorm_binary_f32_f16"); + require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-f32-f16", rmsnorm_f32_f16, token_count_workload(512), + { + { "ggml.rmsnorm_binary_f32.hidden_size", "2048" }, + { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.00001" }, + { "ggml.rmsnorm_binary_f32.op", "2" }, + })); + + const auto & rmsnorm_lowtoken_striped = find_kernel("ggml_rmsnorm_binary_f32_lowtoken_striped"); + require_compiled_kernel(compile_kernel(*jit, "rmsnorm-lowtoken-striped", rmsnorm_lowtoken_striped, + token_count_workload(1), { + { "ggml.rmsnorm_binary_f32.hidden_size", "5120" }, + { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.000001" }, + { "ggml.rmsnorm_binary_f32.op", "2" }, + })); + + const auto & rmsnorm_strided = find_kernel("ggml_rmsnorm_binary_strided_f32"); + require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-strided", rmsnorm_strided, token_count_workload(64), + { + { "ggml.rmsnorm_binary_f32.hidden_size", "256" }, + { "ggml.rmsnorm_binary_f32.input_stride", "288" }, + { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.00001" }, + { "ggml.rmsnorm_binary_f32.op", "2" }, + })); + + 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"); + require_compiled_kernel(compile_kernel(*jit, "prefill-flash-f16", flash_f16, flash_attention_workload(512, 512), + flash_attention_config("32", "8", "64", "0.125"))); + + const auto & swiglu_k16 = find_kernel("ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16"); + require_compiled_kernel(compile_kernel(*jit, "prefill-swiglu-k16", swiglu_k16, token_count_workload(512), + { + { "ggml.mul_mat_swiglu.input_size", "2048" }, + { "ggml.mul_mat_swiglu.output_size", "8192" }, + { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, + { "ggml.mul_mat_swiglu.up_weight_format", "4" }, + { "ggml.mul_mat_swiglu.op", "4" }, + { "ggml.workload.token_capacity", "512" }, + })); + + const auto & get_rows_rmsnorm = find_kernel("ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); + require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-rmsnorm", get_rows_rmsnorm, + { + { "token_count", 512 }, + { "row_count", 128256 }, + { "hidden_size", 2048 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "512" }, + { "ggml.get_rows_f32.hidden_capacity", "2048" }, + { "ggml.get_rows_f32.weight_format", "6" }, + { "ggml.get_rows_rmsnorm.rms_epsilon", "0.00001" }, + })); + + const auto & get_rows_scale = find_kernel("ggml_get_rows_scale_f32"); + require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-scale", get_rows_scale, + { + { "token_count", 64 }, + { "row_count", 122753 }, + { "hidden_size", 2560 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "64" }, + { "ggml.get_rows_f32.hidden_capacity", "2560" }, + { "ggml.get_rows_f32.weight_format", "32" }, + { "ggml.get_rows_scale_f32.scale", "0.177800179" }, + { "ggml.get_rows_scale_f32.bias", "0" }, + })); + require_compiled_kernel(compile_kernel(*jit, "gemma-prefill-get-rows-scale-q5_1", get_rows_scale, + { + { "token_count", 64 }, + { "row_count", 262144 }, + { "hidden_size", 640 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "64" }, + { "ggml.get_rows_f32.hidden_capacity", "640" }, + { "ggml.get_rows_f32.weight_format", "51" }, + { "ggml.get_rows_scale_f32.scale", "25.2982216" }, + { "ggml.get_rows_scale_f32.bias", "0" }, + })); + require_compiled_kernel(compile_kernel(*jit, "gemma-decode-get-rows-scale-q5_1", get_rows_scale, + { + { "token_count", 1 }, + { "row_count", 262144 }, + { "hidden_size", 640 }, + }, + { + { "ggml.get_rows_f32.token_capacity", "1" }, + { "ggml.get_rows_f32.hidden_capacity", "640" }, + { "ggml.get_rows_f32.weight_format", "51" }, + { "ggml.get_rows_scale_f32.scale", "25.2982216" }, + { "ggml.get_rows_scale_f32.bias", "0" }, + })); + + const auto & lfm_dconv3 = find_kernel("llm_ssm_conv_lfm_dconv3_state_f32"); + REQUIRE(lfm_dconv3.bindings.size() == 5); + for (const int64_t tokens : {1, 64}) { + auto ref = compile_kernel(*jit, "lfm-dconv3-" + std::to_string(tokens), lfm_dconv3, {}, { + {"llm.ssm_conv.lfm_dconv3.n_t", std::to_string(tokens)}, + }); + REQUIRE(ref != nullptr); + REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); + auto compiled = ref->take_result(); + REQUIRE(compiled.hsaco_data != nullptr); + REQUIRE(compiled.launch_config.workgroup_size[0] == 256); + REQUIRE(compiled.launch_config.workgroup_count[0] == static_cast(8 * tokens)); + } +} + +static void run_binary_fusion_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 & gather = find_kernel("ggml_gather_add_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-gather-div", gather, + { + { "source_token_count", 2 }, + { "output_token_count", 1 }, + { "hidden_size", 2048 }, + }, + { + { "ggml.gather_add_f32.binary_op", "3" }, + })); + + const auto & scale = find_kernel("ggml_scale_add_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-scale-sub", scale, + { + { "element_count", 128 }, + }, + { + { "ggml.scale_add_f32.ne0", "128" }, + { "ggml.scale_add_f32.ne1", "1" }, + { "ggml.scale_add_f32.ne2", "1" }, + { "ggml.scale_add_f32.input_stride1", "128" }, + { "ggml.scale_add_f32.input_stride2", "128" }, + { "ggml.scale_add_f32.input_stride3", "128" }, + { "ggml.scale_add_f32.input_span", "128" }, + { "ggml.scale_add_f32.residual_stride1", "128" }, + { "ggml.scale_add_f32.residual_stride2", "128" }, + { "ggml.scale_add_f32.residual_stride3", "128" }, + { "ggml.scale_add_f32.residual_span", "128" }, + { "ggml.scale_add_f32.scale", "0.5" }, + { "ggml.scale_add_f32.bias", "0" }, + { "ggml.scale_add_f32.binary_op", "1" }, + { "ggml.scale_add_f32.scaled_lhs", "0" }, + })); + + const auto & rms = find_kernel("ggml_rmsnorm_mul_add_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-rms-div-sub", rms, + { + { "token_count", 1 }, + }, + { + { "ggml.rmsnorm_mul_add_f32.hidden_size", "128" }, + { "ggml.rmsnorm_mul_add_f32.rms_epsilon", "0.000001" }, + { "ggml.rmsnorm_mul_add_f32.normalized_op", "3" }, + { "ggml.rmsnorm_mul_add_f32.output_op", "1" }, + { "ggml.rmsnorm_mul_add_f32.binary_lhs", "0" }, + })); + + const auto & gather_rms = find_kernel("ggml_gather_add_rmsnorm_binary_f32"); + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-gather-rms-mul", gather_rms, + { + { "source_token_count", 2 }, + { "output_token_count", 1 }, + }, + { + { "ggml.gather_add_rmsnorm_binary_f32.hidden_size", "128" }, + { "ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", "0.000001" }, + { "ggml.gather_add_rmsnorm_binary_f32.gather_op", "2" }, + })); + + 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.unary_op", "23" }, + { "llm.ssm_conv.generic.binary_op", "3" }, + { "llm.ssm_conv.generic.binary_lhs", "0" }, + { "llm.ssm_conv.generic.workgroup_size", "256" }, + })); + + const auto & vector_postops = find_kernel("ggml_mul_mat_vector_bias_residual_f32_f32"); + for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { + require_compiled_kernel(compile_kernel( + *jit, "binary-fusion-vector-postops-" + std::to_string(epilogue), vector_postops, token_count_workload(1), + { + { "ggml.matmul.vector_postops.input_size", "2048" }, + { "ggml.matmul.vector_postops.output_size", "2048" }, + { "ggml.matmul.vector_postops.weight_format", "4" }, + { "ggml.matmul.vector_postops.apply_bias", epilogue & 1 ? "1" : "0" }, + { "ggml.matmul.vector_postops.apply_residual", epilogue & 2 ? "1" : "0" }, + { "ggml.workload.token_capacity", "1" }, + })); + } + + const auto compile_postops = [&](const char * kernel_name, const char * key_prefix, int64_t token_count, + int64_t token_capacity) { + const auto & kernel = find_kernel(kernel_name); + for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { + require_compiled_kernel(compile_kernel( + *jit, std::string(key_prefix) + std::to_string(epilogue), kernel, token_count_workload(token_count), + { + { "ggml.mul_mat_postops.input_size", "2048" }, + { "ggml.mul_mat_postops.output_size", "2048" }, + { "ggml.mul_mat_postops.weight_format", "4" }, + { "ggml.mul_mat_postops.epilogue", std::to_string(epilogue) }, + { "ggml.workload.token_capacity", std::to_string(token_capacity) }, + })); + } + }; + compile_postops("ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", "binary-fusion-skinny-postops-", 2, 2); + compile_postops("ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", "binary-fusion-tiled-postops-", 32, 32); + + const auto & pair_postops = find_kernel("ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32"); + for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { + require_compiled_kernel(compile_kernel(*jit, "binary-fusion-pair-postops-" + std::to_string(epilogue), + pair_postops, token_count_workload(32), + { + { "ggml.mul_mat_swiglu.input_size", "2048" }, + { "ggml.mul_mat_swiglu.output_size", "2048" }, + { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, + { "ggml.mul_mat_swiglu.up_weight_format", "4" }, + { "ggml.mul_mat_swiglu.op", "4" }, + { "ggml.mul_mat_swiglu.epilogue", std::to_string(epilogue) }, + { "ggml.workload.token_capacity", "32" }, + })); + } +} + +static void run_loom_jit_materialize_checks() { + const ggml::hrx::KernelDefinition & binary = find_kernel("ggml_binary_f32"); + + HrxTestDevice device; + 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(); +} + +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(); }); + suite.device_case("materialize", [] { run_loom_jit_materialize_checks(); }); +} + +int main(int argc, char ** argv) { + test_runner::Suite suite( + test_runner::Config::with_prefix("HRX Loom JIT test", "hrx-loom-jit", "GGML_HRX_LOOM_JIT_TEST", 2)); + register_hrx_loom_jit_cases(suite); + + const bool has_device = has_hrx_test_device(); + return suite.run(argc, argv, has_device); } diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 0eb83eba8f20..ff745578a939 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -1,19 +1,27 @@ #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" #include "runtime/graph-executor.h" +#include "testing_suite.h" #include +#include +#include #include #include #include #include +#include #include +#include +#include #include #include #include @@ -26,6 +34,21 @@ } \ } while (false) +static void restore_environment_value(const char * name, bool had_value, const std::string & value) { + if (had_value) { + REQUIRE(setenv(name, value.c_str(), 1) == 0); + } else { + REQUIRE(unsetenv(name) == 0); + } +} + +static std::string expected_config_value(float value) { + std::ostringstream out; + out.precision(9); + out << value; + return out.str(); +} + static constexpr float kQwenRmsNormEps = 0.000001f; static constexpr int64_t kQwenFlashHeadSize = 128; static constexpr int64_t kQwenRouterExpertCount = 128; @@ -33,6 +56,9 @@ static constexpr int64_t kQwenRouterRouteCount = 8; static constexpr int64_t kQwenHiddenSize = 2048; static constexpr int64_t kQwenMoeIntermediate = 768; static constexpr int64_t kQwenVocabularyCount = 151936; +static constexpr int64_t kGemmaHiddenSize = 3840; +static constexpr int64_t kGemmaPromptTokenCount = 18; +static constexpr float kGemmaRmsNormEps = 0.000001f; static ggml::hrx::Value make_test_value(ggml::hrx::ValueId id, ggml::hrx::ValueStorageId storage, @@ -77,6 +103,12 @@ static std::vector make_weight(int64_t hidden_size) { return data; } +static const char * expected_plain_rmsnorm_binary_kernel(int64_t hidden_size, int64_t token_count) { + return token_count < 6 && hidden_size >= 2048 && hidden_size <= 5120 ? + "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" : + "loom_libs:ggml_rmsnorm_binary_f32"; +} + static std::vector make_router_input(int64_t hidden_size, int64_t token_count) { std::vector data(hidden_size * token_count); for (int64_t i = 0; i < static_cast(data.size()); ++i) { @@ -103,16 +135,18 @@ static std::vector make_router_logits(int64_t token_count) { return data; } -static std::vector make_flash_query(int64_t token_count) { - std::vector data(kQwenFlashHeadSize * token_count); +static std::vector make_flash_query(int64_t token_count, int64_t head_size = kQwenFlashHeadSize) { + std::vector data(static_cast(head_size) * static_cast(token_count)); for (int64_t i = 0; i < static_cast(data.size()); ++i) { data[i] = static_cast((i % 37) - 18) * 0.01f; } return data; } -static std::vector make_flash_key_value(int64_t token_count, int offset) { - std::vector data(kQwenFlashHeadSize * token_count); +static std::vector make_flash_key_value(int64_t token_count, + int offset, + int64_t head_size = kQwenFlashHeadSize) { + std::vector data(static_cast(head_size) * static_cast(token_count)); for (int64_t i = 0; i < static_cast(data.size()); ++i) { const float value = static_cast(((i + offset) % 31) - 15) * 0.015f; data[i] = ggml_fp32_to_fp16(value); @@ -150,6 +184,25 @@ static std::vector rmsnorm_mul_reference(const std::vector & input return output; } +static std::vector rmsnorm_reference(const std::vector & input, + int64_t hidden_size, + int64_t token_count, + float epsilon) { + std::vector output(input.size()); + for (int64_t token = 0; token < token_count; ++token) { + float sum_squares = 0.0f; + for (int64_t column = 0; column < hidden_size; ++column) { + const float value = input[token * hidden_size + column]; + sum_squares += value * value; + } + const float scale = 1.0f / std::sqrt(sum_squares / static_cast(hidden_size) + epsilon); + for (int64_t column = 0; column < hidden_size; ++column) { + output[token * hidden_size + column] = input[token * hidden_size + column] * scale; + } + } + return output; +} + static std::vector router_projection_reference(const std::vector & input, const std::vector & weight, int64_t hidden_size, @@ -210,19 +263,24 @@ static std::vector flash_attention_reference(const std::vector & const std::vector & value, const std::vector & mask, int64_t query_token_count, - int64_t key_value_token_count) { - std::vector output(query_token_count * kQwenFlashHeadSize); - const float scale = 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)); + int64_t key_value_token_count, + int64_t qk_head_size = kQwenFlashHeadSize, + int64_t value_head_size = -1, + float scale = 0.0f) { + const int64_t actual_value_head_size = value_head_size > 0 ? value_head_size : qk_head_size; + std::vector output(static_cast(query_token_count) * static_cast(actual_value_head_size)); + const float actual_scale = scale == 0.0f ? 1.0f / std::sqrt(static_cast(qk_head_size)) : scale; for (int64_t query_token = 0; query_token < query_token_count; ++query_token) { std::vector scores(key_value_token_count); float max_score = -std::numeric_limits::infinity(); for (int64_t key_token = 0; key_token < key_value_token_count; ++key_token) { float dot = 0.0f; - for (int64_t channel = 0; channel < kQwenFlashHeadSize; ++channel) { - dot += query[query_token * kQwenFlashHeadSize + channel] * - ggml_fp16_to_fp32(key[key_token * kQwenFlashHeadSize + channel]); + for (int64_t channel = 0; channel < qk_head_size; ++channel) { + dot += query[query_token * qk_head_size + channel] * + ggml_fp16_to_fp32(key[key_token * qk_head_size + channel]); } - const float score = dot * scale + ggml_fp16_to_fp32(mask[query_token * key_value_token_count + key_token]); + const float score = + dot * actual_scale + ggml_fp16_to_fp32(mask[query_token * key_value_token_count + key_token]); scores[key_token] = score; max_score = std::max(max_score, score); } @@ -232,13 +290,13 @@ static std::vector flash_attention_reference(const std::vector & score = std::exp(score - max_score); sum += score; } - for (int64_t channel = 0; channel < kQwenFlashHeadSize; ++channel) { + for (int64_t channel = 0; channel < actual_value_head_size; ++channel) { float weighted_sum = 0.0f; for (int64_t key_token = 0; key_token < key_value_token_count; ++key_token) { const float probability = scores[key_token] / sum; - weighted_sum += probability * ggml_fp16_to_fp32(value[key_token * kQwenFlashHeadSize + channel]); + weighted_sum += probability * ggml_fp16_to_fp32(value[key_token * actual_value_head_size + channel]); } - output[query_token * kQwenFlashHeadSize + channel] = weighted_sum; + output[query_token * actual_value_head_size + channel] = weighted_sum; } } return output; @@ -259,9 +317,11 @@ static ggml_tensor * build_qwen_flash_attention_graph(ggml_context * ctx, ggml_tensor * query, ggml_tensor * key, ggml_tensor * value, - ggml_tensor * mask) { - ggml_tensor * output = ggml_flash_attn_ext(ctx, query, key, value, mask, - 1.0f / std::sqrt(static_cast(kQwenFlashHeadSize)), 0.0f, 0.0f); + ggml_tensor * mask, + int64_t head_size = kQwenFlashHeadSize, + float scale = 0.0f) { + const float actual_scale = scale == 0.0f ? 1.0f / std::sqrt(static_cast(head_size)) : scale; + ggml_tensor * output = ggml_flash_attn_ext(ctx, query, key, value, mask, actual_scale, 0.0f, 0.0f); REQUIRE(output != nullptr); return output; } @@ -337,6 +397,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); @@ -387,12 +463,36 @@ 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)) { std::fprintf(stderr, "manual matcher failed for node %zu %s\n", root_index, ggml_op_name(graph.nodes()[root_index].op)); + const ggml::hrx::GraphNode & node = graph.nodes()[root_index]; + for (size_t input_index = 0; input_index < node.inputs.size(); ++input_index) { + const ggml::hrx::Value * input = graph.values().find(node.inputs[input_index]); + if (input != nullptr) { + std::fprintf(stderr, + " input %zu value=%d type=%d ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 + "] nb=[%zu,%zu,%zu,%zu]\n", + input_index, input->id.value, static_cast(input->type), input->ne[0], input->ne[1], + input->ne[2], input->ne[3], input->nb[0], input->nb[1], input->nb[2], input->nb[3]); + } + } + const ggml::hrx::Value * output = graph.values().find(node.output); + if (output != nullptr) { + std::fprintf(stderr, + " output value=%d type=%d ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 + "] nb=[%zu,%zu,%zu,%zu]\n", + output->id.value, static_cast(output->type), output->ne[0], output->ne[1], output->ne[2], + output->ne[3], output->nb[0], output->nb[1], output->nb[2], output->nb[3]); + } + for (const ggml::hrx::CommandPlanAlternateValue & alternate : plan.metadata.alternate_values()) { + std::fprintf(stderr, " alternate graph=%d value=%d type=%d bytes=%zu name=%s\n", + alternate.graph_value.value, alternate.alternate_value.value, static_cast(alternate.type), + alternate.byte_count, alternate.name.c_str()); + } for (const ggml::hrx::DispatchRegistrationAttempt & attempt : diagnostics.attempts) { std::fprintf(stderr, " attempt %s matched=%d\n", attempt.name.c_str(), attempt.matched ? 1 : 0); for (const std::string & error : attempt.errors) { @@ -424,9 +524,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 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; @@ -435,6 +537,11 @@ 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); @@ -449,8 +556,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); @@ -465,6 +605,413 @@ static void run_alternate_value_alias_lookup_checks() { const ggml::hrx::CommandPlanAlternateValue * through_partial_alias = 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) { @@ -476,6 +1023,15 @@ static std::vector make_pattern_f32(size_t element_count, int seed, float return data; } +static std::vector make_gemma_scaled_embedding_input(int64_t hidden_size, int64_t token_count) { + std::vector data = make_pattern_f32(static_cast(hidden_size * token_count), 29, 0.0005f); + const float embedding_scale = std::sqrt(static_cast(hidden_size)); + for (float & value : data) { + value *= embedding_scale; + } + return data; +} + static std::vector make_i32_mod_data(size_t element_count, int32_t modulo) { std::vector data(element_count); for (size_t i = 0; i < element_count; ++i) { @@ -504,9 +1060,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) { @@ -518,6 +1092,88 @@ 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); + std::vector bytes(weights.size() * sizeof(float)); + std::memcpy(bytes.data(), weights.data(), bytes.size()); + return bytes; + } + + if (type == GGML_TYPE_F16) { + (void) seed; + const std::vector weights(static_cast(row_length * row_count), + ggml_fp32_to_fp16(0.00390625f)); + std::vector bytes(weights.size() * sizeof(ggml_fp16_t)); + std::memcpy(bytes.data(), weights.data(), bytes.size()); + return bytes; + } + + if (type == GGML_TYPE_BF16) { + (void) seed; + const std::vector weights(static_cast(row_length * row_count), + ggml_fp32_to_bf16(0.00390625f)); + std::vector bytes(weights.size() * sizeof(ggml_bf16_t)); + std::memcpy(bytes.data(), weights.data(), bytes.size()); + return bytes; + } + + return make_quantized_rows(type, row_length, row_count, seed); +} + static void set_tensor_bytes(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t byte_count) { REQUIRE(tensor != nullptr); REQUIRE(ggml_nbytes(tensor) == byte_count); @@ -562,7 +1218,7 @@ static void require_close(const std::vector & actual, for (size_t i = 0; i < actual.size(); ++i) { const float diff = std::fabs(actual[i] - expected[i]); const float allowed = abs_tolerance + rel_tolerance * std::fabs(expected[i]); - if (diff > allowed) { + if (!std::isfinite(actual[i]) || !std::isfinite(expected[i]) || diff > allowed) { std::fprintf(stderr, "value mismatch at %zu: actual=%g expected=%g diff=%g allowed=%g\n", i, actual[i], expected[i], diff, allowed); std::abort(); @@ -710,19 +1366,23 @@ struct QwenFlashAttentionLayoutGraph { ggml_tensor * output = nullptr; }; -static QwenFlashAttentionLayoutGraph build_qwen_flash_attention_layout_graph(ggml_context * ctx, - int64_t query_token_count, - int64_t key_value_token_count) { +static QwenFlashAttentionLayoutGraph build_qwen_flash_attention_layout_graph( + ggml_context * ctx, + int64_t query_token_count, + int64_t key_value_token_count, + int64_t qk_head_size = kQwenFlashHeadSize, + int64_t value_head_size = kQwenFlashHeadSize, + float scale = 0.0f) { QwenFlashAttentionLayoutGraph graph; constexpr int64_t query_head_count = 32; constexpr int64_t key_value_head_count = 4; ggml_tensor * query_storage = - ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize, query_head_count, query_token_count); + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, qk_head_size, query_head_count, query_token_count); ggml_tensor * key_storage = - ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, key_value_head_count, key_value_token_count); + ggml_new_tensor_3d(ctx, GGML_TYPE_F16, qk_head_size, key_value_head_count, key_value_token_count); ggml_tensor * value_storage = - ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, key_value_head_count, key_value_token_count); + ggml_new_tensor_3d(ctx, GGML_TYPE_F16, value_head_size, key_value_head_count, key_value_token_count); REQUIRE(query_storage != nullptr); REQUIRE(key_storage != nullptr); REQUIRE(value_storage != nullptr); @@ -736,7 +1396,8 @@ static QwenFlashAttentionLayoutGraph build_qwen_flash_attention_layout_graph(ggm REQUIRE(graph.value != nullptr); REQUIRE(graph.mask != nullptr); - graph.output = build_qwen_flash_attention_graph(ctx, graph.query, graph.key, graph.value, graph.mask); + graph.output = + build_qwen_flash_attention_graph(ctx, graph.query, graph.key, graph.value, graph.mask, qk_head_size, scale); return graph; } @@ -909,7 +1570,35 @@ static void run_rmsnorm_support_checks() { REQUIRE(wrong_eps_graph != nullptr); ggml_build_forward_expand(wrong_eps_graph, wrong_eps_output); const ggml::hrx::GraphSupportResult wrong_eps_support = executor.can_execute(*wrong_eps_graph); - REQUIRE(!wrong_eps_support.supported); + REQUIRE(wrong_eps_support.supported); + REQUIRE(wrong_eps_support.status.success()); + require_kernel_subsequence(scheduled_kernel_sequence(wrong_eps_graph), { "loom_libs:ggml_rmsnorm_binary_f32" }); + + ggml_tensor * standalone_output = ggml_rms_norm(ctx, input, 1.0e-5f); + REQUIRE(standalone_output != nullptr); + ggml_cgraph * standalone_graph = ggml_new_graph(ctx); + REQUIRE(standalone_graph != nullptr); + ggml_build_forward_expand(standalone_graph, standalone_output); + const ggml::hrx::GraphSupportResult standalone_support = executor.can_execute(*standalone_graph); + REQUIRE(standalone_support.supported); + REQUIRE(standalone_support.status.success()); + require_kernel_subsequence(scheduled_kernel_sequence(standalone_graph), { "loom_libs:ggml_rmsnorm_f32" }); + + ggml_tensor * q8_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, 1); + ggml_tensor * q8_rhs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); + ggml_tensor * q8_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, kQwenHiddenSize, kQwenVocabularyCount); + REQUIRE(q8_input != nullptr); + REQUIRE(q8_rhs != nullptr); + REQUIRE(q8_weight != nullptr); + ggml_tensor * q8_rms = ggml_rms_norm(ctx, q8_input, 1.0e-5f); + ggml_tensor * q8_binary = ggml_add(ctx, q8_rms, q8_rhs); + ggml_tensor * q8_output = ggml_mul_mat(ctx, q8_weight, q8_binary); + REQUIRE(q8_output != nullptr); + ggml_cgraph * q8_graph = ggml_new_graph(ctx); + REQUIRE(q8_graph != nullptr); + ggml_build_forward_expand(q8_graph, q8_output); + require_kernel_subsequence(scheduled_kernel_sequence(q8_graph), + { "loom_libs:ggml_rmsnorm_binary_f32", "loom_libs:ggml_mul_mat_vector_q6_f32_f32" }); ggml_tensor * wrong_type_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 256, 1); ggml_tensor * wrong_type_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 256); @@ -954,7 +1643,8 @@ static void run_rmsnorm_mul_case(int64_t hidden_size, int64_t token_count) { REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); - require_kernel_subsequence(scheduled_kernel_sequence(graph), { "qwen3_moe:qwen3_moe_rmsnorm_f32" }); + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { expected_plain_rmsnorm_binary_kernel(hidden_size, token_count) }); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -981,435 +1671,2688 @@ static void run_rmsnorm_mul_case(int64_t hidden_size, int64_t token_count) { ggml_backend_free(backend); } -static void run_router_projection_case(int64_t token_count) { - static constexpr int64_t kHiddenSize = 2048; - static constexpr int64_t kExpertCount = 128; +static void run_rmsnorm_binary_add_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_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(backend != nullptr); + constexpr int64_t hidden_size = 256; + constexpr int64_t token_count = 4; + constexpr float epsilon = 1.0e-5f; ggml_init_params params = {}; - params.mem_size = static_cast( - (kHiddenSize * token_count + kHiddenSize * kExpertCount + kExpertCount * token_count) * sizeof(float) * 4 + - 1024 * 1024); - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kHiddenSize, kExpertCount); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kHiddenSize, 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); + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); - require_kernel_subsequence(scheduled_kernel_sequence(graph), - { "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32" }); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * cpu_rhs = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * hrx_rhs = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, hidden_size); + REQUIRE(cpu_input != nullptr); + REQUIRE(cpu_rhs != nullptr); + REQUIRE(hrx_input != nullptr); + REQUIRE(hrx_rhs != nullptr); + + ggml_tensor * cpu_rms = ggml_rms_norm(cpu_ctx, cpu_input, epsilon); + ggml_tensor * hrx_rms = ggml_rms_norm(hrx_ctx, hrx_input, epsilon); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_rms, cpu_rhs); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_rms, hrx_rhs); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); - REQUIRE(buffer != nullptr); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); - const std::vector input_data = make_router_input(kHiddenSize, token_count); - const std::vector weight_data = make_router_weight(kHiddenSize, kExpertCount); - const std::vector expected = - router_projection_reference(input_data, weight_data, kHiddenSize, kExpertCount, token_count); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_binary_f32" }); - ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); - ggml_backend_tensor_set(weight, weight_data.data(), 0, weight_data.size() * sizeof(float)); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); - REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); + const std::vector input = make_pattern_f32(hidden_size * token_count, 17, 0.02f); + const std::vector rhs = make_pattern_f32(hidden_size, 18, 0.03f); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_rhs, hrx_backend, hrx_rhs, rhs.data(), rhs.size() * sizeof(float)); - std::vector actual(expected.size()); - ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); - for (size_t i = 0; i < actual.size(); ++i) { - const float diff = std::fabs(actual[i] - expected[i]); - REQUIRE(diff <= 1.0e-2f); - } + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); - ggml_backend_buffer_free(buffer); - ggml_free(ctx); - ggml_backend_free(backend); + 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_router_top8_case(int64_t token_count) { - ggml_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(backend != nullptr); +static void run_gemma_post_proj_cpu_reference_case(int64_t token_count) { + constexpr int64_t hidden_size = 640; + constexpr float epsilon = 1.0e-6f; + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); ggml_init_params params = {}; - params.mem_size = static_cast(kQwenRouterExpertCount * token_count * sizeof(float) * 16 + 1024 * 1024); - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); - REQUIRE(logits != nullptr); - ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); - require_kernel_subsequence(scheduled_kernel_sequence(graph), { "qwen3_moe:qwen3_moe_router_top8_f32" }); + struct GraphTensors { + ggml_tensor * input; + ggml_tensor * weight; + ggml_tensor * residual; + ggml_tensor * raw; + ggml_tensor * raw_side; + ggml_tensor * output; + }; - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); - REQUIRE(buffer != nullptr); + auto build = [&](ggml_context * ctx) { + GraphTensors g = {}; + g.input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + g.residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + g.raw = ggml_add(ctx, g.input, g.residual); + ggml_tensor * normalized = ggml_rms_norm(ctx, g.raw, epsilon); + ggml_tensor * scaled = ggml_mul(ctx, normalized, g.weight); + g.output = ggml_add(ctx, scaled, g.residual); + g.raw_side = ggml_mul(ctx, g.raw, g.residual); + REQUIRE(g.output != nullptr); + REQUIRE(g.raw_side != nullptr); + return g; + }; + GraphTensors cpu = build(cpu_ctx); + GraphTensors hrx = build(hrx_ctx); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu.output); + ggml_build_forward_expand(cpu_graph, cpu.raw_side); + ggml_build_forward_expand(hrx_graph, hrx.output); + ggml_build_forward_expand(hrx_graph, hrx.raw_side); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_mul_add_f32" }); - const std::vector logits_data = make_router_logits(token_count); - const std::vector expected = router_top8_weights_reference(logits_data, 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); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + const std::vector input = make_pattern_f32(hidden_size * token_count, 37, 0.02f); + const std::vector residual = make_pattern_f32(hidden_size * token_count, 43, 0.04f); + const std::vector weight = make_weight(hidden_size); + set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.residual, hrx_backend, hrx.residual, residual.data(), + residual.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.weight, hrx_backend, hrx.weight, weight.data(), + weight.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), 5.0e-4f); + require_close(get_f32_tensor(hrx_backend, hrx.raw_side), get_f32_tensor(cpu_backend, cpu.raw_side), 1.0e-5f); - ggml_backend_tensor_set(logits, logits_data.data(), 0, logits_data.size() * sizeof(float)); + 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_gemma_post_proj_matcher_negatives() { + constexpr int64_t hidden_size = 640; + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); + REQUIRE(registry != nullptr); + + for (const int variant : { 0, 1, 2, 3, 4 }) { + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 3); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * rhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 3); + ggml_tensor * raw = ggml_add(ctx, input, rhs); + ggml_tensor * residual = variant == 2 ? ggml_view_2d(ctx, raw, hidden_size, 3, raw->nb[1], 0) : + variant == 3 ? ggml_mul(ctx, raw, rhs) : + rhs; + ggml_tensor * rms = ggml_rms_norm(ctx, raw, 1.0e-6f); + ggml_tensor * scaled = variant == 4 ? ggml_add(ctx, rms, weight) : ggml_mul(ctx, rms, weight); + ggml_tensor * output = ggml_add(ctx, scaled, residual); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + if (variant == 1) { + ggml_build_forward_expand(graph, ggml_mul(ctx, scaled, rhs)); + } + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + const size_t rms_index = producer_index_for_tensor(imported.graph, rms); + const size_t raw_index = producer_index_for_tensor(imported.graph, raw); + ggml::hrx::CommandPlan plan; + std::vector covered(imported.graph.nodes().size(), false); + auto kernel = [&]() { + ggml::hrx::DispatchMatch match; + const ggml::hrx::DispatchMatchContext context = { + imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, + plan, next_plan_value(imported.graph, plan), + }; + REQUIRE(registry->match(context, match)); + REQUIRE(match.dispatches.size() == 1); + return kernel_name_for_id(match.dispatches[0].kernel.kernel_id); + }; + REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); + covered[raw_index] = true; + if (variant == 3) { + REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); + covered[producer_index_for_tensor(imported.graph, residual)] = true; + } + if (variant == 2) { + const size_t view_index = producer_index_for_tensor(imported.graph, residual); + covered[view_index] = true; + const ggml::hrx::Value * raw_value = imported.graph.values().find_tensor(raw); + const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(residual); + REQUIRE(raw_value != nullptr); + REQUIRE(residual_value != nullptr); + REQUIRE(raw_value->storage_root == residual_value->storage_root); + } + REQUIRE(kernel() == (variant == 0 || variant == 3 ? "loom_libs:ggml_rmsnorm_mul_add_f32" : + "loom_libs:ggml_rmsnorm_binary_f32")); + ggml_free(ctx); + } +} + +static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, + int64_t head_count, + int64_t token_count, + ggml_unary_op gate_op) { + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t hrx = ggml_backend_hrx_init(0); + REQUIRE(hrx != nullptr); + const int64_t elements = hidden_size * head_count * token_count; + ggml_init_params params = {}; + params.mem_size = static_cast(elements * sizeof(float) * 8 + 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr && hrx_ctx != nullptr); + + struct Graph { + ggml_tensor * input; + ggml_tensor * weight; + ggml_tensor * gate; + ggml_tensor * output; + ggml_cgraph * graph; + }; + + auto build = [&](ggml_context * ctx) { + Graph g; + g.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hidden_size, head_count, token_count); + g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + g.gate = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * head_count, token_count); + ggml_tensor * normalized = build_rmsnorm_mul_graph(ctx, g.input, g.weight); + ggml_tensor * reshaped_gate = ggml_reshape_3d(ctx, g.gate, hidden_size, head_count, token_count); + g.output = ggml_mul(ctx, normalized, ggml_unary(ctx, reshaped_gate, gate_op)); + g.graph = ggml_new_graph(ctx); + ggml_build_forward_expand(g.graph, g.output); + return g; + }; + Graph a = build(cpu_ctx); + Graph b = build(hrx_ctx); + 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); + const auto input = make_pattern_f32(elements, 23, 0.05f); + const auto weight = make_pattern_f32(hidden_size, 11, 0.075f); + const auto gate = make_pattern_f32(elements, 29, 0.125f); + set_tensor_pair_bytes(cpu, a.input, hrx, b.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu, a.weight, hrx, b.weight, weight.data(), weight.size() * sizeof(float)); + set_tensor_pair_bytes(cpu, a.gate, hrx, b.gate, gate.data(), gate.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(cpu, a.graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx, b.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu); + ggml_backend_synchronize(hrx); + const auto actual = get_f32_tensor(hrx, b.output); + const auto cpu_output = get_f32_tensor(cpu, a.output); + if (gate_op == GGML_UNARY_OP_GELU) { + // The CPU GELU route uses an F16 lookup table. Keep the tighter HRX + // check against the defining function instead of inheriting its rounding. + auto expected = rmsnorm_mul_reference(input, weight, hidden_size, head_count * token_count); + for (size_t i = 0; i < expected.size(); ++i) { + const double x = gate[i]; + const double gelu = 0.5 * x * (1.0 + std::tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))); + expected[i] *= static_cast(gelu); + } + require_close(actual, expected, 2.0e-5f); + require_close(cpu_output, expected, 2.0e-5f, 5.0e-4f); + } else { + require_close(actual, cpu_output, 2.0e-5f); + } + + ggml_backend_buffer_free(a_buffer); + ggml_backend_buffer_free(b_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu); + ggml_backend_free(hrx); +} + +static void run_rmsnorm_cpu_reference_case(int64_t hidden_size = 256, + int64_t token_count = 4, + float epsilon = 1.0e-5f) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(hidden_size * token_count * sizeof(float) * 8 + 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, token_count); + REQUIRE(cpu_input != nullptr); + REQUIRE(hrx_input != nullptr); + + ggml_tensor * cpu_output = ggml_rms_norm(cpu_ctx, cpu_input, epsilon); + ggml_tensor * hrx_output = ggml_rms_norm(hrx_ctx, hrx_input, epsilon); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_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); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(hidden_size * token_count, 19, 0.02f); + const std::vector expected = rmsnorm_reference(input, hidden_size, token_count, epsilon); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(cpu_backend, cpu_output), expected, 5.0e-4f); + require_close(get_f32_tensor(hrx_backend, hrx_output), expected, 5.0e-4f); + + 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_strided_rmsnorm_cpu_reference_case(int64_t token_count) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + constexpr int64_t hidden_size = 256; + constexpr int64_t input_stride = 512; + constexpr float epsilon = 1.0e-5f; + constexpr size_t input_offset = 4 * sizeof(float); + const size_t storage_elements = + input_offset / sizeof(float) + static_cast((token_count - 1) * input_stride + hidden_size); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_storage = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * cpu_input = + ggml_view_2d(cpu_ctx, cpu_storage, hidden_size, token_count, input_stride * sizeof(float), input_offset); + ggml_tensor * cpu_output = ggml_rms_norm(cpu_ctx, cpu_input, epsilon); + ggml_tensor * hrx_storage = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * hrx_input = + ggml_view_2d(hrx_ctx, hrx_storage, hidden_size, token_count, input_stride * sizeof(float), input_offset); + ggml_tensor * hrx_output = ggml_rms_norm(hrx_ctx, hrx_input, epsilon); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector rows = make_pattern_f32(static_cast(hidden_size * token_count), 53, 0.015f); + std::vector storage(storage_elements, -19.0f); + for (int64_t token = 0; token < token_count; ++token) { + std::memcpy(storage.data() + input_offset / sizeof(float) + static_cast(token * input_stride), + rows.data() + static_cast(token * hidden_size), + static_cast(hidden_size) * sizeof(float)); + } + const std::vector expected = rmsnorm_reference(rows, hidden_size, token_count, epsilon); + set_tensor_pair_bytes(cpu_backend, cpu_storage, hrx_backend, hrx_storage, storage.data(), + storage.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(cpu_backend, cpu_output), expected, 5.0e-4f); + require_close(get_f32_tensor(hrx_backend, hrx_output), expected, 5.0e-4f); + + 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_strided_rmsnorm_binary_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); + + constexpr int64_t hidden_size = 256; + constexpr int64_t token_count = 64; + constexpr int64_t input_stride = 288; + constexpr float epsilon = 1.0e-5f; + const size_t storage_elements = static_cast((token_count - 1) * input_stride + hidden_size); + + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_storage = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * cpu_input = + ggml_view_2d(cpu_ctx, cpu_storage, hidden_size, token_count, input_stride * sizeof(float), 0); + ggml_tensor * cpu_scale = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * cpu_output = ggml_mul(cpu_ctx, ggml_rms_norm(cpu_ctx, cpu_input, epsilon), cpu_scale); + ggml_tensor * hrx_storage = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * hrx_input = + ggml_view_2d(hrx_ctx, hrx_storage, hidden_size, token_count, input_stride * sizeof(float), 0); + ggml_tensor * hrx_scale = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * hrx_output = ggml_mul(hrx_ctx, ggml_rms_norm(hrx_ctx, hrx_input, epsilon), hrx_scale); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_binary_strided_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector rows = make_pattern_f32(hidden_size * token_count, 54, 0.015f); + const std::vector scale = make_pattern_f32(hidden_size, 55, 0.03f); + std::vector storage(storage_elements, -19.0f); + for (int64_t token = 0; token < token_count; ++token) { + std::memcpy(storage.data() + static_cast(token * input_stride), + rows.data() + static_cast(token * hidden_size), + static_cast(hidden_size) * sizeof(float)); + } + set_tensor_pair_bytes(cpu_backend, cpu_storage, hrx_backend, hrx_storage, storage.data(), + storage.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_scale, hrx_backend, hrx_scale, scale.data(), scale.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + 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_strided_cont_cpu_reference_case(int64_t token_count) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + constexpr int64_t width = 64; + constexpr int64_t row_count = 40; + constexpr int64_t token_stride = 4096; + constexpr size_t view_offset = 8 * sizeof(float); + const size_t storage_elements = + view_offset / sizeof(float) + static_cast((token_count - 1) * token_stride + width * row_count); + + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_storage = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * cpu_view = ggml_view_3d(cpu_ctx, cpu_storage, width, row_count, token_count, width * sizeof(float), + token_stride * sizeof(float), view_offset); + ggml_tensor * cpu_output = ggml_cont(cpu_ctx, cpu_view); + ggml_tensor * hrx_storage = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * hrx_view = ggml_view_3d(hrx_ctx, hrx_storage, width, row_count, token_count, width * sizeof(float), + token_stride * sizeof(float), view_offset); + ggml_tensor * hrx_output = ggml_cont(hrx_ctx, hrx_view); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_copy_strided_source_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector rows = make_pattern_f32(static_cast(width * row_count * token_count), 59, 0.01f); + std::vector storage(storage_elements, -37.0f); + for (int64_t token = 0; token < token_count; ++token) { + std::memcpy(storage.data() + view_offset / sizeof(float) + static_cast(token * token_stride), + rows.data() + static_cast(token * width * row_count), + static_cast(width * row_count) * sizeof(float)); + } + set_tensor_pair_bytes(cpu_backend, cpu_storage, hrx_backend, hrx_storage, storage.data(), + storage.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(cpu_backend, cpu_output), rows, 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_output), rows, 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_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, + int64_t token_count, + float epsilon, + bool adversarial = false) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(hidden_size * token_count * sizeof(float) * 8 + 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * cpu_weight = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, hidden_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * hrx_weight = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, hidden_size); + REQUIRE(cpu_input != nullptr); + REQUIRE(cpu_weight != nullptr); + REQUIRE(hrx_input != nullptr); + REQUIRE(hrx_weight != nullptr); + + ggml_tensor * cpu_output = build_rmsnorm_mul_graph(cpu_ctx, cpu_input, cpu_weight, epsilon); + ggml_tensor * hrx_output = build_rmsnorm_mul_graph(hrx_ctx, hrx_input, hrx_weight, epsilon); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), + { expected_plain_rmsnorm_binary_kernel(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); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + std::vector input = make_pattern_f32(hidden_size * token_count, 23, 0.02f); + if (adversarial) { + for (size_t i = 0; i < input.size(); ++i) { + switch (i % 8) { + case 0: input[i] = 0.0f; break; + case 1: input[i] = 1.0e-6f; break; + case 2: input[i] = -1.0e-6f; break; + case 3: input[i] = 8.0f; break; + case 4: input[i] = -8.0f; break; + case 5: input[i] = 0.125f; break; + case 6: input[i] = -0.25f; break; + default: input[i] = 1.5f; break; + } + } + } + const std::vector weight = make_weight(hidden_size); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), + weight.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + 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_gemma_scaled_rmsnorm_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 = static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount * sizeof(float) * 8 + 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, kGemmaHiddenSize, kGemmaPromptTokenCount); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kGemmaHiddenSize, kGemmaPromptTokenCount); + REQUIRE(cpu_input != nullptr); + REQUIRE(hrx_input != nullptr); + + ggml_tensor * cpu_output = ggml_rms_norm(cpu_ctx, cpu_input, kGemmaRmsNormEps); + ggml_tensor * hrx_output = ggml_rms_norm(hrx_ctx, hrx_input, kGemmaRmsNormEps); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_gemma_scaled_embedding_input(kGemmaHiddenSize, kGemmaPromptTokenCount); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + 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_gemma_scaled_rmsnorm_mul_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 = static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount * sizeof(float) * 8 + 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, kGemmaHiddenSize, kGemmaPromptTokenCount); + ggml_tensor * cpu_weight = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, kGemmaHiddenSize); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, kGemmaHiddenSize, kGemmaPromptTokenCount); + ggml_tensor * hrx_weight = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, kGemmaHiddenSize); + REQUIRE(cpu_input != nullptr); + REQUIRE(cpu_weight != nullptr); + REQUIRE(hrx_input != nullptr); + REQUIRE(hrx_weight != nullptr); + + ggml_tensor * cpu_output = build_rmsnorm_mul_graph(cpu_ctx, cpu_input, cpu_weight, kGemmaRmsNormEps); + ggml_tensor * hrx_output = build_rmsnorm_mul_graph(hrx_ctx, hrx_input, hrx_weight, kGemmaRmsNormEps); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_binary_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_gemma_scaled_embedding_input(kGemmaHiddenSize, kGemmaPromptTokenCount); + const std::vector weight = make_weight(kGemmaHiddenSize); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), + weight.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + 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_scheduled_hrx_scale_rmsnorm_case() { + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + ggml_backend_t cpu_backend = init_cpu_backend(); + REQUIRE(hrx_backend != nullptr); + + ggml_backend_t backends[] = { hrx_backend, cpu_backend }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 2, GGML_DEFAULT_GRAPH_SIZE, false, true); + REQUIRE(sched != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount * sizeof(float) * 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, kGemmaHiddenSize, kGemmaPromptTokenCount); + REQUIRE(input != nullptr); + ggml_backend_sched_set_tensor_backend(sched, input, hrx_backend); + + ggml_tensor * scaled = ggml_scale(ctx, input, std::sqrt(static_cast(kGemmaHiddenSize))); + ggml_tensor * output = ggml_rms_norm(ctx, scaled, kGemmaRmsNormEps); + REQUIRE(scaled != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + 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); + + const std::vector input_data = + make_pattern_f32(static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount), 31, 0.0005f); + std::vector scaled_input = input_data; + const float embedding_scale = std::sqrt(static_cast(kGemmaHiddenSize)); + for (float & value : scaled_input) { + value *= embedding_scale; + } + const std::vector expected = + rmsnorm_reference(scaled_input, kGemmaHiddenSize, kGemmaPromptTokenCount, kGemmaRmsNormEps); + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + ggml_backend_sched_synchronize(sched); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + require_close(actual, expected, 5.0e-4f); + + ggml_free(ctx); + ggml_backend_sched_free(sched); + ggml_backend_free(hrx_backend); + ggml_backend_free(cpu_backend); +} + +static void run_scheduled_hrx_rmsnorm_scale_case() { + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + ggml_backend_t cpu_backend = init_cpu_backend(); + REQUIRE(hrx_backend != nullptr); + + ggml_backend_t backends[] = { hrx_backend, cpu_backend }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 2, GGML_DEFAULT_GRAPH_SIZE, false, true); + REQUIRE(sched != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount * sizeof(float) * 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, kGemmaHiddenSize, kGemmaPromptTokenCount); + REQUIRE(input != nullptr); + ggml_backend_sched_set_tensor_backend(sched, input, hrx_backend); + + ggml_tensor * normalized = ggml_rms_norm(ctx, input, kGemmaRmsNormEps); + ggml_tensor * output = ggml_scale(ctx, normalized, 0.5f); + REQUIRE(normalized != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, normalized) == hrx_backend); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, output) == hrx_backend); + + const std::vector input_data = + make_pattern_f32(static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount), 37, 0.0005f); + std::vector expected = + rmsnorm_reference(input_data, kGemmaHiddenSize, kGemmaPromptTokenCount, kGemmaRmsNormEps); + for (float & value : expected) { + value *= 0.5f; + } + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + ggml_backend_sched_synchronize(sched); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + require_close(actual, expected, 5.0e-4f); + + ggml_free(ctx); + ggml_backend_sched_free(sched); + ggml_backend_free(hrx_backend); + ggml_backend_free(cpu_backend); +} + +static void run_scheduled_hrx_rmsnorm_view_scale_case() { + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + ggml_backend_t cpu_backend = init_cpu_backend(); + REQUIRE(hrx_backend != nullptr); + + ggml_backend_t backends[] = { hrx_backend, cpu_backend }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 2, GGML_DEFAULT_GRAPH_SIZE, false, true); + REQUIRE(sched != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount * sizeof(float) * 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, kGemmaHiddenSize, kGemmaPromptTokenCount); + REQUIRE(input != nullptr); + ggml_backend_sched_set_tensor_backend(sched, input, hrx_backend); + + ggml_tensor * normalized = ggml_rms_norm(ctx, input, kGemmaRmsNormEps); + ggml_tensor * view = ggml_reshape_2d(ctx, normalized, kGemmaHiddenSize, kGemmaPromptTokenCount); + ggml_tensor * output = ggml_scale(ctx, view, 0.25f); + REQUIRE(normalized != nullptr); + REQUIRE(view != nullptr); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, normalized) == hrx_backend); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, output) == hrx_backend); + + const std::vector input_data = + make_pattern_f32(static_cast(kGemmaHiddenSize * kGemmaPromptTokenCount), 41, 0.0005f); + std::vector expected = + rmsnorm_reference(input_data, kGemmaHiddenSize, kGemmaPromptTokenCount, kGemmaRmsNormEps); + for (float & value : expected) { + value *= 0.25f; + } + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + ggml_backend_sched_synchronize(sched); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + require_close(actual, expected, 5.0e-4f); + + ggml_free(ctx); + ggml_backend_sched_free(sched); + ggml_backend_free(hrx_backend); + ggml_backend_free(cpu_backend); +} + +static void run_external_view_input_rmsnorm_case() { + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + static constexpr int64_t kHeadSize = 256; + static constexpr int64_t kHeadCount = 16; + const int64_t row_count = kHeadCount * kGemmaPromptTokenCount; + + ggml_init_params params = {}; + params.mem_size = static_cast(kHeadSize * row_count * sizeof(float) * 4 + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * root = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kHeadSize, kHeadCount, kGemmaPromptTokenCount, 1); + REQUIRE(root != nullptr); + ggml_tensor * input_view = ggml_view_4d(ctx, root, kHeadSize, kHeadCount, kGemmaPromptTokenCount, 1, root->nb[1], + root->nb[2], root->nb[3], 0); + REQUIRE(input_view != nullptr); + ggml_tensor * output = ggml_rms_norm(ctx, input_view, kGemmaRmsNormEps); + REQUIRE(output != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_graph_add_node(graph, output); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_rmsnorm_f32" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, hrx_backend); + REQUIRE(buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(kHeadSize * row_count), 29, 0.0025f); + const std::vector expected = rmsnorm_reference(input, kHeadSize, row_count, kGemmaRmsNormEps); + + ggml_backend_tensor_set(root, input.data(), 0, input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(hrx_backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, output), expected, 5.0e-4f); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(hrx_backend); +} + +static void run_split_local_view_alias_import_case() { + static constexpr int64_t kHeadSize = 256; + static constexpr int64_t kKeyValueHeads = 8; + static constexpr int64_t kKeyValueHidden = kHeadSize * kKeyValueHeads; + static constexpr int64_t kRootHeads = kKeyValueHeads + 1; + const size_t split_offset = static_cast(kHeadSize * sizeof(float)); + + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = + static_cast(kHeadSize * kRootHeads * kGemmaPromptTokenCount * sizeof(float) * 4 + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * root = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kHeadSize, kRootHeads, kGemmaPromptTokenCount, 1); + REQUIRE(root != nullptr); + ggml_tensor * split_input = ggml_view_4d(ctx, root, kHeadSize, kKeyValueHeads, kGemmaPromptTokenCount, 1, + root->nb[1], root->nb[2], root->nb[3], split_offset); + REQUIRE(split_input != nullptr); + ggml_tensor * flattened = + ggml_view_2d(ctx, split_input, kKeyValueHidden, kGemmaPromptTokenCount, split_input->nb[2], 0); + REQUIRE(flattened != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_graph_add_node(graph, flattened); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(ggml::hrx::is_layout_alias_node(imported.graph, imported.graph.nodes().front())); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); + REQUIRE(scheduler.plan().dispatches.empty()); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, hrx_backend); + REQUIRE(buffer != nullptr); + + const size_t root_element_count = static_cast(kHeadSize * kRootHeads * kGemmaPromptTokenCount); + const std::vector root_data = make_pattern_f32(root_element_count, 37, 0.001f); + ggml_backend_tensor_set(root, root_data.data(), 0, root_data.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(hrx_backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(hrx_backend); + + std::vector actual(static_cast(kKeyValueHidden * kGemmaPromptTokenCount)); + ggml_backend_tensor_get(flattened, actual.data(), 0, actual.size() * sizeof(float)); + const float * expected = root_data.data() + split_offset / sizeof(float); + require_close(actual, std::vector(expected, expected + actual.size()), 0.0f); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(hrx_backend); +} + +static void run_router_projection_case(int64_t token_count) { + static constexpr int64_t kHiddenSize = 2048; + static constexpr int64_t kExpertCount = 128; + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast( + (kHiddenSize * token_count + kHiddenSize * kExpertCount + kExpertCount * token_count) * sizeof(float) * 4 + + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kHiddenSize, kExpertCount); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kHiddenSize, 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); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), + { "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector input_data = make_router_input(kHiddenSize, token_count); + const std::vector weight_data = make_router_weight(kHiddenSize, kExpertCount); + const std::vector expected = + router_projection_reference(input_data, weight_data, kHiddenSize, kExpertCount, token_count); + + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); + ggml_backend_tensor_set(weight, weight_data.data(), 0, weight_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 1.0e-2f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_router_top8_case(int64_t token_count) { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(kQwenRouterExpertCount * token_count * sizeof(float) * 16 + 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenRouterExpertCount, token_count); + REQUIRE(logits != nullptr); + ggml_tensor * output = build_qwen_router_top8_graph(ctx, logits); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + require_kernel_subsequence(scheduled_kernel_sequence(graph), { "qwen3_moe:qwen3_moe_router_top8_f32" }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector logits_data = make_router_logits(token_count); + const std::vector expected = router_top8_weights_reference(logits_data, token_count); + + ggml_backend_tensor_set(logits, logits_data.data(), 0, logits_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 1.0e-5f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_qwen_flash_attention_case() { + static constexpr int64_t kQueryTokenCount = 2; + static constexpr int64_t kKeyValueTokenCount = 4; + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize, kQueryTokenCount, 1); + ggml_tensor * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, kKeyValueTokenCount, 1); + ggml_tensor * value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, kKeyValueTokenCount, 1); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kKeyValueTokenCount, kQueryTokenCount); + REQUIRE(query != nullptr); + REQUIRE(key != nullptr); + REQUIRE(value != nullptr); + REQUIRE(mask != nullptr); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, query, key, value, mask); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + + 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; + const std::string original_env_value = had_original_env ? original_env : ""; + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + ggml_tensor * common_query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize, 16, 1); + ggml_tensor * common_key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, 16, 1); + ggml_tensor * common_value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, 16, 1); + ggml_tensor * common_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 16, 16); + REQUIRE(common_query != nullptr); + REQUIRE(common_key != nullptr); + REQUIRE(common_value != nullptr); + REQUIRE(common_mask != nullptr); + ggml_tensor * common_output = + build_qwen_flash_attention_graph(ctx, common_query, common_key, common_value, common_mask); + ggml_cgraph * common_graph = ggml_new_graph(ctx); + REQUIRE(common_graph != nullptr); + ggml_build_forward_expand(common_graph, common_output); + require_kernel_subsequence(scheduled_kernel_sequence(common_graph), + { "loom_libs:ggml_flash_attention_f32_f16_wmma" }); + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + + const std::vector query_data = make_flash_query(kQueryTokenCount); + const std::vector key_data = make_flash_key_value(kKeyValueTokenCount, 3); + const std::vector value_data = make_flash_key_value(kKeyValueTokenCount, 11); + const std::vector mask_data = make_flash_mask(kQueryTokenCount, kKeyValueTokenCount); + const std::vector expected = + flash_attention_reference(query_data, key_data, value_data, mask_data, kQueryTokenCount, kKeyValueTokenCount); + + ggml_backend_tensor_set(query, query_data.data(), 0, query_data.size() * sizeof(float)); + ggml_backend_tensor_set(key, key_data.data(), 0, key_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(value, value_data.data(), 0, value_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 5.0e-2f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_common_flash_attention_cpu_reference_case(int64_t head_size, + int64_t query_token_count, + int64_t key_value_token_count, + float scale, + const char * expected_kernel, + int64_t value_head_size = -1) { + if (value_head_size < 0) { + value_head_size = head_size; + } + 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; + const std::string original_env_value = had_original_env ? original_env : ""; + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + 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 * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, query_token_count, 1); + ggml_tensor * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, head_size, key_value_token_count, 1); + ggml_tensor * value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, value_head_size, key_value_token_count, 1); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_token_count, query_token_count); + REQUIRE(query != nullptr); + REQUIRE(key != nullptr); + REQUIRE(value != nullptr); + REQUIRE(mask != nullptr); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, query, key, value, mask, head_size, scale); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + 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); + + const std::vector query_data = make_flash_query(query_token_count, head_size); + const std::vector key_data = make_flash_key_value(key_value_token_count, 3, head_size); + const std::vector value_data = make_flash_key_value(key_value_token_count, 11, value_head_size); + const std::vector mask_data = make_flash_mask(query_token_count, key_value_token_count); + const std::vector expected = + flash_attention_reference(query_data, key_data, value_data, mask_data, query_token_count, key_value_token_count, + head_size, value_head_size, scale); + + ggml_backend_tensor_set(query, query_data.data(), 0, query_data.size() * sizeof(float)); + ggml_backend_tensor_set(key, key_data.data(), 0, key_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(value, value_data.data(), 0, value_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); + + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 5.0e-2f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); +} + +static void run_asymmetric_flash_attention_cpu_reference_case() { + static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; + const char * original_env = std::getenv(kDisableQwenDispatchEnv); + const bool had_original_env = original_env != nullptr; + const std::string original_env_value = had_original_env ? original_env : ""; + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + + constexpr int64_t query_token_count = 23; + constexpr int64_t key_value_token_count = 256; + constexpr int64_t qk_head_size = 96; + constexpr int64_t value_head_size = 64; + const float scale = 1.0f / std::sqrt(static_cast(qk_head_size)); + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + + 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 * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, qk_head_size, query_token_count, 1); + ggml_tensor * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, qk_head_size, key_value_token_count, 1); + ggml_tensor * value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, value_head_size, key_value_token_count, 1); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, key_value_token_count, query_token_count); + REQUIRE(query != nullptr); + REQUIRE(key != nullptr); + REQUIRE(value != nullptr); + REQUIRE(mask != nullptr); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, query, key, value, mask, qk_head_size, scale); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + 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); + + const std::vector query_data = make_flash_query(query_token_count, qk_head_size); + const std::vector key_data = make_flash_key_value(key_value_token_count, 3, qk_head_size); + const std::vector value_data = make_flash_key_value(key_value_token_count, 11, value_head_size); + const std::vector mask_data = make_flash_mask(query_token_count, key_value_token_count); + const std::vector expected = + flash_attention_reference(query_data, key_data, value_data, mask_data, query_token_count, key_value_token_count, + qk_head_size, value_head_size, scale); + + ggml_backend_tensor_set(query, query_data.data(), 0, query_data.size() * sizeof(float)); + ggml_backend_tensor_set(key, key_data.data(), 0, key_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(value, value_data.data(), 0, value_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(backend); - std::vector actual(expected.size()); - ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); - for (size_t i = 0; i < actual.size(); ++i) { - const float diff = std::fabs(actual[i] - expected[i]); - REQUIRE(diff <= 1.0e-5f); + std::vector actual(expected.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + REQUIRE(diff <= 5.0e-2f); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); +} + +static void run_qwen_full_cache_prefill_flash_attention_scheduling_case(int64_t token_count) { + constexpr int64_t kQueryHeadCount = 32; + constexpr int64_t kKeyValueHeadCount = 4; + constexpr int64_t kFullCacheRowCount = 40960; + const size_t active_mask_byte_count = + static_cast(token_count) * static_cast(token_count) * sizeof(ggml_fp16_t); + + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + AttentionPostprocessGraph attention = + build_attention_postprocess_graph(ctx, token_count, kQueryHeadCount, kKeyValueHeadCount, kFullCacheRowCount); + ggml_tensor * flash_output = append_qwen_full_cache_flash_attention_consumer( + ctx, attention, token_count, kQueryHeadCount, kKeyValueHeadCount, kFullCacheRowCount); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, attention.key_output); + ggml_build_forward_expand(graph, attention.value_output); + ggml_build_forward_expand(graph, flash_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); + REQUIRE(registry != nullptr); + + std::vector covered_nodes(imported.graph.nodes().size(), false); + ggml::hrx::CommandPlan plan; + ggml::hrx::DispatchMatch postprocess_match; + match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, + producer_index_for_tensor(imported.graph, attention.query_reshape), postprocess_match); + ggml::hrx::DispatchMatch flash_match; + match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, + producer_index_for_tensor(imported.graph, flash_output), flash_match); + + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.attention_mask); + REQUIRE(mask_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * compact_mask = + ggml::hrx::find_alternate_value(plan, mask_value->id, GGML_TYPE_F16, active_mask_byte_count); + REQUIRE(compact_mask != nullptr); + + const ggml::hrx::Dispatch * metadata_dispatch = nullptr; + for (const ggml::hrx::Dispatch & dispatch : plan.initialization_dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:qwen_attention_metadata") { + metadata_dispatch = &dispatch; + } + } + REQUIRE(metadata_dispatch != nullptr); + REQUIRE(metadata_dispatch->kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(metadata_dispatch->kernel.integer_parameters.at("context_capacity") == token_count); + REQUIRE(metadata_dispatch->bindings.size() == 5); + REQUIRE(metadata_dispatch->bindings[4].value == compact_mask->alternate_value); + REQUIRE(metadata_dispatch->bindings[4].length == active_mask_byte_count); + + const ggml::hrx::Dispatch * flash_dispatch = nullptr; + for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_flash_attention_f32_f16_wmma") { + flash_dispatch = &dispatch; + } + } + REQUIRE(flash_dispatch != nullptr); + REQUIRE(flash_dispatch->kernel.integer_parameters.at("query_token_count") == token_count); + REQUIRE(flash_dispatch->kernel.integer_parameters.at("key_value_token_count") == token_count); + REQUIRE(flash_dispatch->bindings.size() == 6); + REQUIRE(flash_dispatch->kernel.compile_parameters.at("ggml.flash_attention.apply_gate") == "0"); + REQUIRE(flash_dispatch->bindings[3].value == compact_mask->alternate_value); + REQUIRE(flash_dispatch->bindings[3].length == active_mask_byte_count); + + ggml_free(ctx); +} + +static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_token_count, + int64_t key_value_token_count, + bool use_common_dispatch = false, + int64_t qk_head_size = kQwenFlashHeadSize, + int64_t value_head_size = kQwenFlashHeadSize, + float scale = 0.0f) { + 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; + const std::string original_env_value = had_original_env ? original_env : ""; + if (use_common_dispatch) { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + } else { + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + + constexpr int64_t query_head_count = 32; + constexpr int64_t key_value_head_count = 4; + const int64_t query_hidden_size = query_head_count * qk_head_size; + 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; + 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); + const size_t partial_output_bytes = + static_cast(key_value_head_count * key_value_blocks * 16 * value_head_size) * sizeof(ggml_fp16_t); + + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + 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, consumer_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == static_cast(query_token_count + 1)); + REQUIRE(plan.transients.size() == 4); + REQUIRE(plan.completion_counter_requests.size() == 1); + + REQUIRE(plan.transients[0].size == partial_scalar_bytes); + REQUIRE(plan.transients[1].size == partial_scalar_bytes); + REQUIRE(plan.transients[2].size == partial_output_bytes); + 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 * 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(publication_value->id, GGML_TYPE_Q8_1, q8_output_bytes); + REQUIRE(q8_alternate != nullptr); + REQUIRE(q8_alternate->alternate_value == plan.transients[3].value); + + const size_t query_row_bytes = static_cast(query_hidden_size) * sizeof(float); + const size_t mask_row_bytes = static_cast(key_value_token_count) * sizeof(ggml_fp16_t); + const size_t output_row_bytes = static_cast(output_hidden_size) * sizeof(float); + for (int64_t row = 0; row < query_token_count; ++row) { + 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)); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.key_value_head_count") == + std::to_string(key_value_head_count)); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.qk_head_size") == + std::to_string(qk_head_size)); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.value_head_size") == + std::to_string(value_head_size)); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.attention_scale") == + expected_config_value(scale == 0.0f ? 1.0f / std::sqrt(static_cast(qk_head_size)) : scale)); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.decode.key_value_token_capacity") == + std::to_string(key_value_capacity)); + REQUIRE(dispatch.bindings.size() == 10); + REQUIRE(dispatch.bindings[0].value == query_value->id); + REQUIRE(dispatch.bindings[0].offset == static_cast(row) * query_value->nb[1]); + REQUIRE(dispatch.bindings[0].length == query_row_bytes); + REQUIRE(dispatch.bindings[3].value == mask_value->id); + REQUIRE(dispatch.bindings[3].offset == static_cast(row) * mask_value->nb[1]); + REQUIRE(dispatch.bindings[3].length == mask_row_bytes); + REQUIRE(dispatch.bindings[8].value == output_value->id); + REQUIRE(dispatch.bindings[8].offset == static_cast(row) * output_value->nb[2]); + REQUIRE(dispatch.bindings[8].length == output_row_bytes); + REQUIRE(dispatch.bindings[9].value == q8_alternate->alternate_value); + REQUIRE(dispatch.bindings[9].offset == static_cast(row) * q8_row_bytes); + REQUIRE(dispatch.bindings[9].length == q8_row_bytes); + } + + ggml_free(ctx); + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); +} + +static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool include_get_rows_selectors) { + constexpr int64_t query_token_count = 1; + constexpr int64_t key_value_token_count = 512; + constexpr int64_t attention_hidden_size = 4096; + const size_t attention_q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, attention_hidden_size); + const size_t next_q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, kQwenHiddenSize); + + ggml_init_params params = {}; + params.mem_size = 128 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + QwenFlashAttentionLayoutGraph attention = + build_qwen_flash_attention_layout_graph(ctx, query_token_count, key_value_token_count); + ggml_tensor * attention_output = ggml_reshape_2d(ctx, attention.output, attention_hidden_size, query_token_count); + ggml_tensor * output_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, attention_hidden_size, kQwenHiddenSize); + ggml_tensor * projection = ggml_mul_mat(ctx, output_weight, attention_output); + ggml_tensor * residual_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, query_token_count); + ggml_tensor * selected_projection = projection; + ggml_tensor * selected_residual = residual_input; + ggml_tensor * row_indices = nullptr; + if (include_get_rows_selectors) { + row_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + selected_projection = ggml_get_rows(ctx, projection, row_indices); + selected_residual = ggml_get_rows(ctx, residual_input, row_indices); + REQUIRE(row_indices != nullptr); + REQUIRE(selected_projection != nullptr); + REQUIRE(selected_residual != nullptr); + } + ggml_tensor * residual = ggml_add(ctx, selected_projection, selected_residual); + 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); + REQUIRE(residual_input != nullptr); + REQUIRE(residual != nullptr); + 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, next_output); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == 3); + REQUIRE(plan.completion_counter_requests.size() == 2); + + 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); + REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.input_size") == + std::to_string(attention_hidden_size)); + REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_size") == + std::to_string(kQwenHiddenSize)); + 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 * 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); + REQUIRE(flash_output_value != nullptr); + REQUIRE(residual_input_value != nullptr); + REQUIRE(residual_value != nullptr); + REQUIRE(normalized_value != nullptr); + const ggml::hrx::CommandPlanAlternateValue * attention_q8 = + 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(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); + REQUIRE(projection_dispatch.bindings[2].value == residual_value->id); + REQUIRE(projection_dispatch.bindings[4].value == normalized_value->id); + REQUIRE(projection_dispatch.bindings[6].value == next_q8->alternate_value); + + const ggml::hrx::Value * aliased_residual = imported.graph.values().find(residual_value->id); + REQUIRE(aliased_residual != nullptr); + REQUIRE(aliased_residual->alias_source == residual_input_value->id); + + ggml_free(ctx); +} + +static void run_add_f32_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 cpu_params = {}; + cpu_params.mem_size = 256 * 1024; + cpu_params.no_alloc = true; + ggml_init_params hrx_params = cpu_params; + ggml_context * cpu_ctx = ggml_init(cpu_params); + ggml_context * hrx_ctx = ggml_init(hrx_params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t element_count = 257; + ggml_tensor * cpu_a = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_b = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_a, cpu_b); + ggml_tensor * hrx_a = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_b = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_a, hrx_b); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_binary_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector a = make_pattern_f32(element_count, 1, 0.125f); + const std::vector b = make_pattern_f32(element_count, 2, 0.25f); + set_tensor_pair_bytes(cpu_backend, cpu_a, hrx_backend, hrx_a, a.data(), a.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_b, hrx_backend, hrx_b, b.data(), b.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 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_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); + 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 hidden_size = 2048; + constexpr int64_t token_count = 2; + ggml_tensor * cpu_source = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, hidden_size, 3, token_count); + ggml_tensor * hrx_source = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, hidden_size, 3, token_count); + REQUIRE(cpu_source != nullptr); + REQUIRE(hrx_source != nullptr); + ggml_tensor * cpu_lhs = ggml_view_2d(cpu_ctx, cpu_source, hidden_size, token_count, cpu_source->nb[2], 0); + ggml_tensor * cpu_rhs = + ggml_view_2d(cpu_ctx, cpu_source, hidden_size, token_count, cpu_source->nb[2], 2 * cpu_source->nb[1]); + ggml_tensor * hrx_lhs = ggml_view_2d(hrx_ctx, hrx_source, hidden_size, token_count, hrx_source->nb[2], 0); + ggml_tensor * hrx_rhs = + ggml_view_2d(hrx_ctx, hrx_source, hidden_size, token_count, hrx_source->nb[2], 2 * hrx_source->nb[1]); + REQUIRE(cpu_lhs != nullptr); + REQUIRE(cpu_rhs != nullptr); + REQUIRE(hrx_lhs != nullptr); + REQUIRE(hrx_rhs != nullptr); + ggml_tensor * cpu_output = ggml_mul(cpu_ctx, cpu_lhs, cpu_rhs); + ggml_tensor * hrx_output = ggml_mul(hrx_ctx, hrx_lhs, hrx_rhs); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_binary_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector source = make_pattern_f32(hidden_size * 3 * token_count, 19, 0.015625f); + set_tensor_pair_bytes(cpu_backend, cpu_source, hrx_backend, hrx_source, source.data(), + source.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 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_glu_split_f32_cpu_reference_case(ggml_glu_op glu_op) { + 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 = 256 * 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 hidden_size = 256; + constexpr int64_t token_count = 3; + ggml_tensor * cpu_gate = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * cpu_up = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * cpu_output = ggml_glu_split(cpu_ctx, cpu_gate, cpu_up, glu_op); + ggml_tensor * hrx_gate = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * hrx_up = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * hrx_output = ggml_glu_split(hrx_ctx, hrx_gate, hrx_up, glu_op); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_binary_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector gate = make_pattern_f32(hidden_size * token_count, 21, 0.02f); + const std::vector up = make_pattern_f32(hidden_size * token_count, 22, 0.03f); + set_tensor_pair_bytes(cpu_backend, cpu_gate, hrx_backend, hrx_gate, gate.data(), gate.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_up, hrx_backend, hrx_up, up.data(), up.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + 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_glu_packed_f32_cpu_reference_case(ggml_glu_op glu_op, + bool swapped, + int64_t hidden_size = 256, + int64_t token_count = 3) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(16 * 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 2 * hidden_size, token_count); + ggml_tensor * cpu_output = ggml_glu(cpu_ctx, cpu_input, glu_op, swapped); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 2 * hidden_size, token_count); + ggml_tensor * hrx_output = ggml_glu(hrx_ctx, hrx_input, glu_op, swapped); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_binary_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = + make_pattern_f32(static_cast(2 * hidden_size * token_count), swapped ? 24 : 23, 0.02f); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + + 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_gather_add_f32_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 hidden_size = 256; + constexpr int64_t source_token_count = 11; + constexpr int64_t output_token_count = 5; + ggml_tensor * cpu_a = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * cpu_b = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * cpu_rows_a = ggml_get_rows(cpu_ctx, cpu_a, cpu_ids); + ggml_tensor * cpu_rows_b = ggml_get_rows(cpu_ctx, cpu_b, cpu_ids); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_rows_a, cpu_rows_b); + + ggml_tensor * hrx_a = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * hrx_b = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, source_token_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, output_token_count); + ggml_tensor * hrx_rows_a = ggml_get_rows(hrx_ctx, hrx_a, hrx_ids); + ggml_tensor * hrx_rows_b = ggml_get_rows(hrx_ctx, hrx_b, hrx_ids); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_rows_a, hrx_rows_b); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "hrx:ggml_gather_add_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector a = make_pattern_f32(hidden_size * source_token_count, 3, 0.05f); + const std::vector b = make_pattern_f32(hidden_size * source_token_count, 4, 0.075f); + const std::vector ids = { 9, 3, 7, 1, 5 }; + set_tensor_pair_bytes(cpu_backend, cpu_a, hrx_backend, hrx_a, a.data(), a.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_b, hrx_backend, hrx_b, b.data(), b.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 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_scale_add_f32_cpu_reference_case(bool strided_residual) { + 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 width = 257; + constexpr int64_t token_count = 3; + constexpr int64_t residual_width = 320; + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, width, token_count); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, width, token_count); + ggml_tensor * cpu_residual_storage = + ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, strided_residual ? residual_width : width, token_count); + ggml_tensor * hrx_residual_storage = + ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, strided_residual ? residual_width : width, token_count); + ggml_tensor * cpu_residual = strided_residual ? ggml_view_2d(cpu_ctx, cpu_residual_storage, width, token_count, + cpu_residual_storage->nb[1], 0) : + cpu_residual_storage; + ggml_tensor * hrx_residual = strided_residual ? ggml_view_2d(hrx_ctx, hrx_residual_storage, width, token_count, + hrx_residual_storage->nb[1], 0) : + hrx_residual_storage; + ggml_tensor * cpu_scaled = ggml_scale_bias(cpu_ctx, cpu_input, 0.177800179f, -0.125f); + ggml_tensor * hrx_scaled = ggml_scale_bias(hrx_ctx, hrx_input, 0.177800179f, -0.125f); + ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_scaled, cpu_residual); + ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_scaled, hrx_residual); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_scale_add_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(width * token_count, 31, 0.03125f); + const std::vector residual = + make_pattern_f32((strided_residual ? residual_width : width) * token_count, 37, 0.046875f); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_residual_storage, hrx_backend, hrx_residual_storage, residual.data(), + residual.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 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_get_rows_f32_cpu_reference_case(ggml_type weight_type, + int64_t hidden_size = kQwenHiddenSize, + int64_t vocabulary_count = 64, + int64_t token_count = 7) { + 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 = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, hidden_size, vocabulary_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_output = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, hidden_size, vocabulary_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_output = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_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); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector weight = make_matmul_weight_bytes(weight_type, hidden_size, vocabulary_count, 5); + std::vector ids(static_cast(token_count)); + for (int64_t i = 0; i < token_count; ++i) { + ids[static_cast(i)] = static_cast((i * 17 + 3) % vocabulary_count); } + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); - ggml_backend_buffer_free(buffer); - ggml_free(ctx); - ggml_backend_free(backend); -} + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); -static void run_qwen_flash_attention_case() { - static constexpr int64_t kQueryTokenCount = 2; - static constexpr int64_t kKeyValueTokenCount = 4; + 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); +} - ggml_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(backend != nullptr); +static void run_get_rows_q8_1_zero_weight_case() { + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); ggml_init_params params = {}; - params.mem_size = 2 * 1024 * 1024; + params.mem_size = 4 * 1024 * 1024; params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(hrx_ctx != nullptr); - ggml_tensor * query = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kQwenFlashHeadSize, kQueryTokenCount, 1); - ggml_tensor * key = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, kKeyValueTokenCount, 1); - ggml_tensor * value = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, kQwenFlashHeadSize, kKeyValueTokenCount, 1); - ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kKeyValueTokenCount, kQueryTokenCount); - REQUIRE(query != nullptr); - REQUIRE(key != nullptr); - REQUIRE(value != nullptr); - REQUIRE(mask != nullptr); - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, query, key, value, mask); + constexpr int64_t vocabulary_count = 64; + constexpr int64_t hidden_size = kQwenHiddenSize; + constexpr int64_t token_count = 7; + ggml_tensor * weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q8_1, hidden_size, vocabulary_count); + ggml_tensor * ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * output = ggml_get_rows(hrx_ctx, weight, ids); + REQUIRE(output != nullptr); - ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_cgraph * graph = ggml_new_graph(hrx_ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); - require_kernel_subsequence(scheduled_kernel_sequence(graph), - { "qwen3_moe:qwen3_moe_flash_attention_f32_f16_wmma" }); + require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_get_rows_f32" }); - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(buffer != nullptr); - const std::vector query_data = make_flash_query(kQueryTokenCount); - const std::vector key_data = make_flash_key_value(kKeyValueTokenCount, 3); - const std::vector value_data = make_flash_key_value(kKeyValueTokenCount, 11); - const std::vector mask_data = make_flash_mask(kQueryTokenCount, kKeyValueTokenCount); - const std::vector expected = - flash_attention_reference(query_data, key_data, value_data, mask_data, kQueryTokenCount, kKeyValueTokenCount); - - ggml_backend_tensor_set(query, query_data.data(), 0, query_data.size() * sizeof(float)); - ggml_backend_tensor_set(key, key_data.data(), 0, key_data.size() * sizeof(ggml_fp16_t)); - ggml_backend_tensor_set(value, value_data.data(), 0, value_data.size() * sizeof(ggml_fp16_t)); - ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); - - REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); + const std::vector weight_bytes(ggml_nbytes(weight), 0); + const std::vector id_values = { 3, 17, 29, 41, 53, 7, 19 }; + set_tensor_bytes(hrx_backend, weight, weight_bytes.data(), weight_bytes.size()); + set_tensor_bytes(hrx_backend, ids, id_values.data(), id_values.size() * sizeof(int32_t)); - std::vector actual(expected.size()); - ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); - for (size_t i = 0; i < actual.size(); ++i) { - const float diff = std::fabs(actual[i] - expected[i]); - REQUIRE(diff <= 5.0e-2f); - } + REQUIRE(ggml_backend_graph_compute(hrx_backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, output), + std::vector(static_cast(hidden_size * token_count), 0.0f), 0.0f, 0.0f); ggml_backend_buffer_free(buffer); - ggml_free(ctx); - ggml_backend_free(backend); + ggml_free(hrx_ctx); + ggml_backend_free(hrx_backend); } -static void run_qwen_full_cache_prefill_flash_attention_scheduling_case(int64_t token_count) { - constexpr int64_t kQueryHeadCount = 32; - constexpr int64_t kKeyValueHeadCount = 4; - constexpr int64_t kFullCacheRowCount = 40960; - const size_t active_mask_byte_count = - static_cast(token_count) * static_cast(token_count) * sizeof(ggml_fp16_t); +static void run_get_rows_scale_f32_cpu_reference_case(ggml_type weight_type = GGML_TYPE_F32, + int64_t hidden_size = 256, + int64_t row_count = 8, + int64_t token_count = 6) { + 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 = 128 * 1024 * 1024; + params.mem_size = 4 * 1024 * 1024; params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); - AttentionPostprocessGraph attention = - build_attention_postprocess_graph(ctx, token_count, kQueryHeadCount, kKeyValueHeadCount, kFullCacheRowCount); - ggml_tensor * flash_output = append_qwen_full_cache_flash_attention_consumer( - ctx, attention, token_count, kQueryHeadCount, kKeyValueHeadCount, kFullCacheRowCount); + const float scale = weight_type == GGML_TYPE_Q5_1 ? 25.2982216f : 0.177800179f; + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, hidden_size, row_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); + ggml_tensor * cpu_output = ggml_scale(cpu_ctx, cpu_rows, scale); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, hidden_size, row_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * hrx_output = ggml_scale(hrx_ctx, hrx_rows, scale); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); + ggml_set_output(cpu_rows); + ggml_set_output(hrx_rows); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, attention.key_output); - ggml_build_forward_expand(graph, attention.value_output); - ggml_build_forward_expand(graph, flash_output); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_get_rows_scale_f32" }); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); - REQUIRE(registry != nullptr); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::CommandPlan plan; - ggml::hrx::DispatchMatch postprocess_match; - match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, - producer_index_for_tensor(imported.graph, attention.query_reshape), postprocess_match); - ggml::hrx::DispatchMatch flash_match; - match_dispatch_at_index(imported.graph, *registry, plan, covered_nodes, - producer_index_for_tensor(imported.graph, flash_output), flash_match); + const std::vector weight = make_matmul_weight_bytes(weight_type, hidden_size, row_count, 9); + const std::vector ids = { 0, static_cast(row_count - 1), 2, + 2, static_cast(row_count - 2), 1 }; + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); - const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.attention_mask); - REQUIRE(mask_value != nullptr); - const ggml::hrx::CommandPlanAlternateValue * compact_mask = - ggml::hrx::find_alternate_value(plan, mask_value->id, GGML_TYPE_F16, active_mask_byte_count); - REQUIRE(compact_mask != nullptr); + 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); + const float tolerance = weight_type == GGML_TYPE_F32 ? 0.0f : 5.0e-4f; + require_close(get_f32_tensor(hrx_backend, hrx_rows), get_f32_tensor(cpu_backend, cpu_rows), tolerance); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), tolerance); - const ggml::hrx::Dispatch * metadata_dispatch = nullptr; - for (const ggml::hrx::Dispatch & dispatch : plan.initialization_dispatches) { - if (kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:qwen_attention_metadata") { - metadata_dispatch = &dispatch; - } + 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 auto allocate_test_weight(ggml_backend_t backend, ggml_tensor * weight) { + const auto buft = ggml_backend_get_default_buffer_type(backend); + const size_t alignment = ggml_backend_buft_get_alignment(buft); + const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, weight), alignment); + auto buffer = std::unique_ptr( + ggml_backend_buft_alloc_buffer(buft, size), ggml_backend_buffer_free); + REQUIRE(buffer != nullptr); + ggml_tallocr allocator = ggml_tallocr_new(buffer.get()); + REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(buffer.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + return buffer; +} + +static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size, + int64_t input_size = kQwenHiddenSize, + bool normalize_input = false, + bool square_output = false, + ggml_type cache_type = GGML_TYPE_COUNT) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(32 * 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_norm_w = normalize_input ? ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, input_size) : nullptr; + ggml_tensor * hrx_norm_w = normalize_input ? ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, input_size) : nullptr; + ggml_tensor * cpu_rhs = normalize_input ? build_rmsnorm_mul_graph(cpu_ctx, cpu_input, cpu_norm_w) : cpu_input; + ggml_tensor * hrx_rhs = normalize_input ? build_rmsnorm_mul_graph(hrx_ctx, hrx_input, hrx_norm_w) : hrx_input; + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_rhs); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_rhs); + if (square_output) { + cpu_output = ggml_sqr(cpu_ctx, cpu_output); + hrx_output = ggml_sqr(hrx_ctx, hrx_output); } - REQUIRE(metadata_dispatch != nullptr); - REQUIRE(metadata_dispatch->kernel.integer_parameters.at("token_count") == token_count); - REQUIRE(metadata_dispatch->kernel.integer_parameters.at("context_capacity") == token_count); - REQUIRE(metadata_dispatch->bindings.size() == 5); - REQUIRE(metadata_dispatch->bindings[4].value == compact_mask->alternate_value); - REQUIRE(metadata_dispatch->bindings[4].length == active_mask_byte_count); + ggml_tensor * cpu_indices = nullptr; + ggml_tensor * hrx_indices = nullptr; + ggml_tensor * cpu_cache = nullptr; + ggml_tensor * hrx_cache = nullptr; + if (cache_type != GGML_TYPE_COUNT) { + REQUIRE(token_count == 1); + cpu_indices = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I64, token_count); + hrx_indices = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I64, token_count); + cpu_cache = ggml_new_tensor_2d(cpu_ctx, cache_type, output_size, 17); + hrx_cache = ggml_new_tensor_2d(hrx_ctx, cache_type, output_size, 17); + cpu_output = ggml_set_rows(cpu_ctx, cpu_cache, cpu_output, cpu_indices); + hrx_output = ggml_set_rows(hrx_ctx, hrx_cache, hrx_output, hrx_indices); + } + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); - const ggml::hrx::Dispatch * flash_dispatch = nullptr; - for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { - if (kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_flash_attention_f32_f16_wmma") { - flash_dispatch = &dispatch; + auto hrx_weight_buffer = allocate_test_weight(hrx_backend, hrx_weight); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + const auto sequence = scheduled_kernel_sequence(hrx_graph); + require_kernel_subsequence(sequence, { expected_kernel }); + const bool packed_q4_input = weight_type == GGML_TYPE_Q4_K; + const bool iq4_xs_weight = weight_type == GGML_TYPE_IQ4_XS; + const bool prefill_tokens = token_count >= 256; + const bool packed_iq4_xs_input = iq4_xs_weight && prefill_tokens; + 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_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_publish") == sequence.end()); + } + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector weight = make_matmul_weight_bytes(weight_type, input_size, output_size, 6); + const bool weight_is_f16 = weight_type == GGML_TYPE_F16; + const bool weight_is_bf16 = weight_type == GGML_TYPE_BF16; + const bool weight_is_f32 = weight_type == GGML_TYPE_F32; + const bool weight_is_dense_float = weight_is_f16 || weight_is_bf16 || weight_is_f32; + std::vector input = weight_is_dense_float ? + std::vector(static_cast(input_size * token_count), 0.00390625f) : + make_pattern_f32(input_size * token_count, 7, 0.01f); + if (cache_type != GGML_TYPE_COUNT) { + for (size_t i = 0; i < input.size(); ++i) { + const int value = i % 32 == 0 ? 127 : static_cast((i * 29) % 255) - 127; + input[i] = static_cast(value) * 0.0078125f; } } - REQUIRE(flash_dispatch != nullptr); - REQUIRE(flash_dispatch->kernel.integer_parameters.at("query_token_count") == token_count); - REQUIRE(flash_dispatch->kernel.integer_parameters.at("key_value_token_count") == token_count); - REQUIRE(flash_dispatch->bindings.size() == 5); - REQUIRE(flash_dispatch->bindings[3].value == compact_mask->alternate_value); - REQUIRE(flash_dispatch->bindings[3].length == active_mask_byte_count); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + if (normalize_input) { + const std::vector norm_w = make_weight(input_size); + set_tensor_pair_bytes(cpu_backend, cpu_norm_w, hrx_backend, hrx_norm_w, norm_w.data(), + norm_w.size() * sizeof(float)); + } + if (cache_type != GGML_TYPE_COUNT) { + const int64_t index = 7; + set_tensor_pair_bytes(cpu_backend, cpu_indices, hrx_backend, hrx_indices, &index, sizeof(index)); + const std::vector cache(ggml_nbytes(cpu_cache), 0); + set_tensor_pair_bytes(cpu_backend, cpu_cache, hrx_backend, hrx_cache, cache.data(), cache.size()); + } - ggml_free(ctx); + 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); + const float abs_tolerance = cache_type != GGML_TYPE_COUNT ? 2.0e-3f : 1.0f; + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), abs_tolerance, + cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); + + 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_qwen_decode_split_flash_attention_scheduling_case(int64_t query_token_count, - int64_t key_value_token_count) { - constexpr int64_t query_head_count = 32; - constexpr int64_t key_value_head_count = 4; - constexpr int64_t hidden_size = query_head_count * kQwenFlashHeadSize; - const size_t q8_row_bytes = ggml_row_size(GGML_TYPE_Q8_1, 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; - 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); - const size_t partial_output_bytes = - static_cast(key_value_head_count * key_value_blocks * 16 * kQwenFlashHeadSize) * sizeof(ggml_fp16_t); +static void run_prefill_value_cache_repeated_indices_case() { + static constexpr const char * disable_env = "GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"; + const char * old_env = std::getenv(disable_env); + const bool had_env = old_env != nullptr; + const std::string old_value = had_env ? old_env : ""; + REQUIRE(unsetenv(disable_env) == 0); + + constexpr int64_t input_size = 1536; + constexpr int64_t output_size = 256; + constexpr int64_t token_count = 64; + constexpr int64_t vocabulary_count = 128; + constexpr int64_t cache_row_count = 512; + + ggml_backend_t fused_backend = ggml_backend_hrx_init(0); + ggml_backend_t split_backend = ggml_backend_hrx_init(0); + REQUIRE(fused_backend != nullptr); + REQUIRE(split_backend != nullptr); ggml_init_params params = {}; - params.mem_size = 128 * 1024 * 1024; + params.mem_size = static_cast(32 * 1024 * 1024); params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); + ggml_context * fused_ctx = ggml_init(params); + ggml_context * split_ctx = ggml_init(params); + REQUIRE(fused_ctx != nullptr); + REQUIRE(split_ctx != nullptr); + + auto build_graph = [](ggml_context * ctx) { + ggml_tensor * embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, input_size, vocabulary_count); + ggml_tensor * tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, input_size); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, input_size, output_size); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, output_size, cache_row_count); + ggml_tensor * rows = ggml_get_rows(ctx, embedding, tokens); + ggml_tensor * input = build_rmsnorm_mul_graph(ctx, rows, norm); + ggml_tensor * projected = ggml_mul_mat(ctx, weight, input); + ggml_tensor * output = ggml_set_rows(ctx, cache, projected, indices); + return std::array{ embedding, tokens, norm, weight, indices, cache, output }; + }; - QwenFlashAttentionLayoutGraph attention = - build_qwen_flash_attention_layout_graph(ctx, query_token_count, key_value_token_count); + const auto fused = build_graph(fused_ctx); + const auto split = build_graph(split_ctx); + auto fused_embedding_buffer = allocate_test_weight(fused_backend, fused[0]); + auto fused_projection_buffer = allocate_test_weight(fused_backend, fused[3]); + auto split_embedding_buffer = allocate_test_weight(split_backend, split[0]); + auto split_projection_buffer = allocate_test_weight(split_backend, split[3]); + + ggml_cgraph * fused_graph = ggml_new_graph(fused_ctx); + ggml_cgraph * split_graph = ggml_new_graph(split_ctx); + REQUIRE(fused_graph != nullptr); + REQUIRE(split_graph != nullptr); + ggml_build_forward_expand(fused_graph, fused[6]); + ggml_build_forward_expand(split_graph, split[6]); + require_kernel_subsequence(scheduled_kernel_sequence(fused_graph), + { "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", + "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32" }); + 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_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); + REQUIRE(fused_buffer != nullptr); + REQUIRE(split_buffer != nullptr); + + const std::vector embedding = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, vocabulary_count, 5); + const std::vector projection = make_matmul_weight_bytes(GGML_TYPE_Q6_K, input_size, output_size, 11); + const std::vector norm = make_weight(input_size); + const std::vector tokens(static_cast(token_count), 3); + std::vector indices(static_cast(token_count)); + for (int64_t token = 0; token < token_count; ++token) { + indices[static_cast(token)] = (token * 37 + 11) % 31; + } + const std::vector cache(static_cast(output_size * cache_row_count), ggml_fp32_to_fp16(-2.5f)); + + set_tensor_pair_bytes(fused_backend, fused[0], split_backend, split[0], embedding.data(), embedding.size()); + set_tensor_pair_bytes(fused_backend, fused[1], split_backend, split[1], tokens.data(), + tokens.size() * sizeof(int32_t)); + set_tensor_pair_bytes(fused_backend, fused[2], split_backend, split[2], norm.data(), norm.size() * sizeof(float)); + set_tensor_pair_bytes(fused_backend, fused[3], split_backend, split[3], projection.data(), projection.size()); + set_tensor_pair_bytes(fused_backend, fused[4], split_backend, split[4], indices.data(), + indices.size() * sizeof(int64_t)); + set_tensor_pair_bytes(fused_backend, fused[5], split_backend, split[5], cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + + REQUIRE(unsetenv(disable_env) == 0); + REQUIRE(ggml_backend_graph_compute(fused_backend, fused_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(fused_backend); + REQUIRE(setenv(disable_env, "1", 1) == 0); + REQUIRE(ggml_backend_graph_compute(split_backend, split_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(split_backend); + require_close(get_f32_tensor(fused_backend, fused[6]), get_f32_tensor(split_backend, split[6]), 0.0f); + + ggml_backend_buffer_free(fused_buffer); + ggml_backend_buffer_free(split_buffer); + ggml_free(fused_ctx); + ggml_free(split_ctx); + ggml_backend_free(fused_backend); + ggml_backend_free(split_backend); + restore_environment_value(disable_env, had_env, old_value); +} + +static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size, + int64_t input_size = kQwenHiddenSize) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(32 * 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); - ggml_cgraph * cgraph = ggml_new_graph(ctx); - REQUIRE(cgraph != nullptr); - ggml_build_forward_expand(cgraph, attention.output); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_matmul = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_tensor * cpu_output = ggml_sqr(cpu_ctx, cpu_matmul); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_matmul = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * hrx_output = ggml_sqr(hrx_ctx, hrx_matmul); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - ggml::hrx::DispatchScheduleDiagnostics diagnostics; - REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == static_cast(query_token_count)); - REQUIRE(plan.transients.size() == 4); - REQUIRE(plan.completion_counter_requests.size() == 1); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); - REQUIRE(plan.transients[0].size == partial_scalar_bytes); - REQUIRE(plan.transients[1].size == partial_scalar_bytes); - REQUIRE(plan.transients[2].size == partial_output_bytes); - REQUIRE(plan.transients[3].size == q8_output_bytes); - REQUIRE(plan.completion_counter_requests[0].count == key_value_head_count); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); - 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); - REQUIRE(query_value != nullptr); - REQUIRE(mask_value != nullptr); - REQUIRE(output_value != nullptr); - const ggml::hrx::CommandPlanAlternateValue * q8_alternate = - plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_Q8_1, q8_output_bytes); - REQUIRE(q8_alternate != nullptr); - REQUIRE(q8_alternate->alternate_value == plan.transients[3].value); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); - const size_t query_row_bytes = static_cast(hidden_size) * sizeof(float); - const size_t mask_row_bytes = static_cast(key_value_token_count) * sizeof(ggml_fp16_t); - const size_t output_row_bytes = query_row_bytes; - for (int64_t row = 0; row < query_token_count; ++row) { - const ggml::hrx::Dispatch & dispatch = plan.dispatches[static_cast(row)]; - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == - "qwen3_moe:qwen3_moe_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.kernel.compile_parameters.at("qwen3_moe.attention.key_value_token_capacity") == - std::to_string(key_value_capacity)); - REQUIRE(dispatch.bindings.size() == 10); - REQUIRE(dispatch.bindings[0].value == query_value->id); - REQUIRE(dispatch.bindings[0].offset == static_cast(row) * query_value->nb[1]); - REQUIRE(dispatch.bindings[0].length == query_row_bytes); - REQUIRE(dispatch.bindings[3].value == mask_value->id); - REQUIRE(dispatch.bindings[3].offset == static_cast(row) * mask_value->nb[1]); - REQUIRE(dispatch.bindings[3].length == mask_row_bytes); - REQUIRE(dispatch.bindings[8].value == output_value->id); - REQUIRE(dispatch.bindings[8].offset == static_cast(row) * output_value->nb[2]); - REQUIRE(dispatch.bindings[8].length == output_row_bytes); - REQUIRE(dispatch.bindings[9].value == q8_alternate->alternate_value); - REQUIRE(dispatch.bindings[9].offset == static_cast(row) * q8_row_bytes); - REQUIRE(dispatch.bindings[9].length == q8_row_bytes); - } + const std::vector weight(ggml_nbytes(cpu_weight), 0); + const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); - ggml_free(ctx); + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0e-1f, 3.0e-2f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); } -static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool include_get_rows_selectors) { - constexpr int64_t query_token_count = 1; - constexpr int64_t key_value_token_count = 512; - constexpr int64_t attention_hidden_size = 4096; - const size_t attention_q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, attention_hidden_size); - const size_t next_q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, kQwenHiddenSize); +static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_type, + ggml_type up_weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size, + ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, + int64_t input_size = kQwenHiddenSize) { + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); ggml_init_params params = {}; - params.mem_size = 128 * 1024 * 1024; + params.mem_size = static_cast(48 * 1024 * 1024); params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); - QwenFlashAttentionLayoutGraph attention = - build_qwen_flash_attention_layout_graph(ctx, query_token_count, key_value_token_count); - ggml_tensor * attention_output = ggml_reshape_2d(ctx, attention.output, attention_hidden_size, query_token_count); - ggml_tensor * output_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, attention_hidden_size, kQwenHiddenSize); - ggml_tensor * projection = ggml_mul_mat(ctx, output_weight, attention_output); - ggml_tensor * residual_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenHiddenSize, query_token_count); - ggml_tensor * selected_projection = projection; - ggml_tensor * selected_residual = residual_input; - ggml_tensor * row_indices = nullptr; - if (include_get_rows_selectors) { - row_indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - selected_projection = ggml_get_rows(ctx, projection, row_indices); - selected_residual = ggml_get_rows(ctx, residual_input, row_indices); - REQUIRE(row_indices != nullptr); - REQUIRE(selected_projection != nullptr); - REQUIRE(selected_residual != nullptr); - } - ggml_tensor * residual = ggml_add(ctx, selected_projection, selected_residual); - 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); - REQUIRE(attention_output != nullptr); - REQUIRE(output_weight != nullptr); - REQUIRE(projection != nullptr); - REQUIRE(residual_input != nullptr); - REQUIRE(residual != nullptr); - REQUIRE(norm_weight != nullptr); - REQUIRE(rms != nullptr); - REQUIRE(normalized != nullptr); + ggml_tensor * cpu_gate_weight = ggml_new_tensor_2d(cpu_ctx, gate_weight_type, input_size, output_size); + ggml_tensor * cpu_up_weight = ggml_new_tensor_2d(cpu_ctx, up_weight_type, input_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_gate = ggml_mul_mat(cpu_ctx, cpu_gate_weight, cpu_input); + ggml_tensor * cpu_up = ggml_mul_mat(cpu_ctx, cpu_up_weight, cpu_input); + ggml_tensor * cpu_output = ggml_glu_split(cpu_ctx, cpu_gate, cpu_up, glu_op); + ggml_tensor * hrx_gate_weight = ggml_new_tensor_2d(hrx_ctx, gate_weight_type, input_size, output_size); + ggml_tensor * hrx_up_weight = ggml_new_tensor_2d(hrx_ctx, up_weight_type, input_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_gate = ggml_mul_mat(hrx_ctx, hrx_gate_weight, hrx_input); + ggml_tensor * hrx_up = ggml_mul_mat(hrx_ctx, hrx_up_weight, hrx_input); + ggml_tensor * hrx_output = ggml_glu_split(hrx_ctx, hrx_gate, hrx_up, glu_op); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); - ggml_cgraph * cgraph = ggml_new_graph(ctx); - REQUIRE(cgraph != nullptr); - ggml_build_forward_expand(cgraph, normalized); + auto hrx_gate_buffer = allocate_test_weight(hrx_backend, hrx_gate_weight); + auto hrx_up_buffer = allocate_test_weight(hrx_backend, hrx_up_weight); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - ggml::hrx::DispatchScheduleDiagnostics diagnostics; - REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 2); - REQUIRE(plan.completion_counter_requests.size() == 2); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); - const ggml::hrx::Dispatch & projection_dispatch = plan.dispatches.back(); - 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); - REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.input_size") == - std::to_string(attention_hidden_size)); - REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_size") == - std::to_string(kQwenHiddenSize)); - REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_accumulation") == "1"); - REQUIRE(projection_dispatch.bindings.size() == 7); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); - 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); - REQUIRE(flash_output_value != nullptr); - REQUIRE(residual_input_value != nullptr); - 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); - const ggml::hrx::CommandPlanAlternateValue * next_q8 = - ggml::hrx::find_alternate_value(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); - REQUIRE(projection_dispatch.bindings[2].value == residual_value->id); - REQUIRE(projection_dispatch.bindings[4].value == normalized_value->id); - REQUIRE(projection_dispatch.bindings[6].value == next_q8->alternate_value); + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); - const ggml::hrx::Value * aliased_residual = imported.graph.values().find(residual_value->id); - REQUIRE(aliased_residual != nullptr); - REQUIRE(aliased_residual->alias_source == residual_input_value->id); + const std::vector gate_weight = make_matmul_weight_bytes(gate_weight_type, input_size, output_size, 6); + const std::vector up_weight = make_matmul_weight_bytes(up_weight_type, input_size, output_size, 11); + const std::vector input(static_cast(input_size * token_count), 0.00390625f); + set_tensor_pair_bytes(cpu_backend, cpu_gate_weight, hrx_backend, hrx_gate_weight, gate_weight.data(), + gate_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_up_weight, hrx_backend, hrx_up_weight, up_weight.data(), up_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); - ggml_free(ctx); + REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu_backend); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0e-4f, 1.0e-3f); + + 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_add_f32_cpu_reference_case() { +static void run_dense_matmul_packed_glu_cpu_reference_case(ggml_type weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size, + ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, + bool swapped = false, + int64_t input_size = 256) { 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 cpu_params = {}; - cpu_params.mem_size = 256 * 1024; - cpu_params.no_alloc = true; - ggml_init_params hrx_params = cpu_params; - ggml_context * cpu_ctx = ggml_init(cpu_params); - ggml_context * hrx_ctx = ggml_init(hrx_params); + ggml_init_params params = {}; + params.mem_size = static_cast(64 * 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t element_count = 257; - ggml_tensor * cpu_a = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * cpu_b = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_a, cpu_b); - ggml_tensor * hrx_a = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * hrx_b = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_a, hrx_b); + ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, 2 * output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_packed = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_tensor * cpu_output = ggml_glu(cpu_ctx, cpu_packed, glu_op, swapped); + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, 2 * output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_packed = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * hrx_output = ggml_glu(hrx_ctx, hrx_packed, glu_op, swapped); REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); @@ -1420,23 +4363,23 @@ static void run_add_f32_cpu_reference_case() { ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "qwen3_moe:ggml_add_f32" }); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - const std::vector a = make_pattern_f32(element_count, 1, 0.125f); - const std::vector b = make_pattern_f32(element_count, 2, 0.25f); - set_tensor_pair_bytes(cpu_backend, cpu_a, hrx_backend, hrx_a, a.data(), a.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu_b, hrx_backend, hrx_b, b.data(), b.size() * sizeof(float)); + const std::vector weight = make_matmul_weight_bytes(weight_type, input_size, 2 * output_size, 15); + const std::vector input(static_cast(input_size * token_count), 0.00390625f); + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0f, 3.0e-2f); ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -1446,38 +4389,42 @@ static void run_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_gather_add_f32_cpu_reference_case() { +static void run_dense_matmul_binary_cpu_reference_case(ggml_type lhs_weight_type, + ggml_type rhs_weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size) { 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.mem_size = static_cast(48 * 1024 * 1024); params.no_alloc = true; ggml_context * cpu_ctx = ggml_init(params); ggml_context * hrx_ctx = ggml_init(params); REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t hidden_size = 256; - constexpr int64_t source_token_count = 11; - constexpr int64_t output_token_count = 5; - ggml_tensor * cpu_a = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, source_token_count); - ggml_tensor * cpu_b = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, source_token_count); - ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, output_token_count); - ggml_tensor * cpu_rows_a = ggml_get_rows(cpu_ctx, cpu_a, cpu_ids); - ggml_tensor * cpu_rows_b = ggml_get_rows(cpu_ctx, cpu_b, cpu_ids); - ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_rows_a, cpu_rows_b); - - ggml_tensor * hrx_a = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, source_token_count); - ggml_tensor * hrx_b = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, source_token_count); - ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, output_token_count); - ggml_tensor * hrx_rows_a = ggml_get_rows(hrx_ctx, hrx_a, hrx_ids); - ggml_tensor * hrx_rows_b = ggml_get_rows(hrx_ctx, hrx_b, hrx_ids); - ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_rows_a, hrx_rows_b); + constexpr int64_t input_size = kQwenHiddenSize; + ggml_tensor * cpu_lhs_weight = ggml_new_tensor_2d(cpu_ctx, lhs_weight_type, input_size, output_size); + ggml_tensor * cpu_rhs_weight = ggml_new_tensor_2d(cpu_ctx, rhs_weight_type, input_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_lhs = ggml_mul_mat(cpu_ctx, cpu_lhs_weight, cpu_input); + ggml_tensor * cpu_rhs = ggml_mul_mat(cpu_ctx, cpu_rhs_weight, cpu_input); + ggml_tensor * cpu_output = ggml_sub(cpu_ctx, cpu_lhs, cpu_rhs); + ggml_tensor * hrx_lhs_weight = ggml_new_tensor_2d(hrx_ctx, lhs_weight_type, input_size, output_size); + ggml_tensor * hrx_rhs_weight = ggml_new_tensor_2d(hrx_ctx, rhs_weight_type, input_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_lhs = ggml_mul_mat(hrx_ctx, hrx_lhs_weight, hrx_input); + ggml_tensor * hrx_rhs = ggml_mul_mat(hrx_ctx, hrx_rhs_weight, hrx_input); + ggml_tensor * hrx_output = ggml_sub(hrx_ctx, hrx_lhs, hrx_rhs); REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); + auto hrx_lhs_buffer = allocate_test_weight(hrx_backend, hrx_lhs_weight); + auto hrx_rhs_buffer = allocate_test_weight(hrx_backend, hrx_rhs_weight); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); REQUIRE(cpu_graph != nullptr); @@ -1485,25 +4432,27 @@ static void run_gather_add_f32_cpu_reference_case() { ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "qwen3_moe:ggml_gather_add_f32" }); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - const std::vector a = make_pattern_f32(hidden_size * source_token_count, 3, 0.05f); - const std::vector b = make_pattern_f32(hidden_size * source_token_count, 4, 0.075f); - const std::vector ids = { 9, 3, 7, 1, 5 }; - set_tensor_pair_bytes(cpu_backend, cpu_a, hrx_backend, hrx_a, a.data(), a.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu_b, hrx_backend, hrx_b, b.data(), b.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); + const std::vector lhs_weight = make_matmul_weight_bytes(lhs_weight_type, input_size, output_size, 6); + const std::vector rhs_weight = make_matmul_weight_bytes(rhs_weight_type, input_size, output_size, 11); + const std::vector input(static_cast(input_size * token_count), 0.00390625f); + set_tensor_pair_bytes(cpu_backend, cpu_lhs_weight, hrx_backend, hrx_lhs_weight, lhs_weight.data(), + lhs_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_rhs_weight, hrx_backend, hrx_rhs_weight, rhs_weight.data(), + rhs_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0e-4f, 1.0e-3f); ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -1513,55 +4462,103 @@ static void run_gather_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_token_embedding_q4k_cpu_reference_case() { +static void run_vector_q8_publish_cpu_reference_case() { + constexpr int64_t input_size = 256; + constexpr int64_t intermediate_size = 256; + constexpr int64_t output_size = 128; + constexpr int64_t token_count = 1; + ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; + params.mem_size = static_cast(32 * 1024 * 1024); params.no_alloc = true; ggml_context * cpu_ctx = ggml_init(params); ggml_context * hrx_ctx = ggml_init(params); REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t vocabulary_count = 64; - constexpr int64_t hidden_size = kQwenHiddenSize; - constexpr int64_t token_count = 7; - ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, hidden_size, vocabulary_count); - ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); - ggml_tensor * cpu_output = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); - ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, hidden_size, vocabulary_count); - ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); - ggml_tensor * hrx_output = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); + ggml_tensor * cpu_first_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, intermediate_size); + ggml_tensor * cpu_second_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, intermediate_size, output_size); + ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * cpu_hidden = ggml_mul_mat(cpu_ctx, cpu_first_weight, cpu_input); + ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_second_weight, cpu_hidden); + ggml_tensor * hrx_first_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, intermediate_size); + ggml_tensor * hrx_second_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, intermediate_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_hidden = ggml_mul_mat(hrx_ctx, hrx_first_weight, hrx_input); + ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_second_weight, hrx_hidden); REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_output); - ggml_build_forward_expand(hrx_graph, hrx_output); + auto hrx_first_weight_buffer = allocate_test_weight(hrx_backend, hrx_first_weight); + auto hrx_second_weight_buffer = allocate_test_weight(hrx_backend, hrx_second_weight); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), + { "loom_libs:ggml_mul_mat_vector_f32_f32", "qwen3_moe:ggml_quantize_q8_1_x4_f32", + "loom_libs:ggml_mul_mat_vector_f32_f32" }); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*hrx_graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + ggml::hrx::DispatchScheduleDiagnostics diagnostics; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); + const ggml::hrx::CommandPlan & plan = scheduler.plan(); + REQUIRE(plan.valid()); + + const ggml::hrx::Value * hidden_value = imported.graph.values().find_tensor(hrx_hidden); + REQUIRE(hidden_value != nullptr); + const size_t q8_byte_count = static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, intermediate_size); + const ggml::hrx::CommandPlanAlternateValue * q8_alternate = + plan.metadata.find_alternate_value(hidden_value->id, GGML_TYPE_Q8_1, q8_byte_count); + REQUIRE(q8_alternate != nullptr); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "qwen3_moe:qwen_token_embedding_q4k" }); + const ggml::hrx::Dispatch * second_vector = nullptr; + int vector_count = 0; + for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { + if (kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_vector_f32_f32") { + ++vector_count; + if (vector_count == 2) { + second_vector = &dispatch; + break; + } + } + } + REQUIRE(second_vector != nullptr); + REQUIRE(second_vector->bindings.size() == 4); + REQUIRE(second_vector->bindings[0].value == q8_alternate->alternate_value); + REQUIRE(second_vector->kernel.compile_parameters.at("ggml.matmul.vector.activation_format") == + std::to_string(GGML_TYPE_Q8_1)); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - const std::vector weight = make_quantized_rows(GGML_TYPE_Q4_K, hidden_size, vocabulary_count, 5); - const std::vector ids = { 3, 17, 29, 41, 53, 7, 19 }; - set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); + const std::vector first_weight = make_matmul_weight_bytes(GGML_TYPE_F32, input_size, intermediate_size, 6); + const std::vector second_weight = + make_matmul_weight_bytes(GGML_TYPE_Q4_K, intermediate_size, output_size, 11); + const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); + set_tensor_pair_bytes(cpu_backend, cpu_first_weight, hrx_backend, hrx_first_weight, first_weight.data(), + first_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_second_weight, hrx_backend, hrx_second_weight, second_weight.data(), + second_weight.size()); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0f, 5.0e-2f); ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -1571,29 +4568,69 @@ static void run_token_embedding_q4k_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, - const char * expected_kernel, - int64_t token_count, - int64_t output_size) { +static void run_dense_matmul_postops_cpu_reference_case(ggml_type weight_type, + const char * expected_kernel, + int64_t token_count, + int64_t output_size, + bool include_bias, + bool include_residual, + bool include_next_rmsnorm, + int64_t input_size = kQwenHiddenSize) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); ggml_init_params params = {}; - params.mem_size = static_cast(32 * 1024 * 1024); + params.mem_size = static_cast(96 * 1024 * 1024); params.no_alloc = true; ggml_context * cpu_ctx = ggml_init(params); ggml_context * hrx_ctx = ggml_init(params); REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t input_size = kQwenHiddenSize; ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); - ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); - ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * cpu_bias = include_bias ? ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, output_size) : nullptr; + ggml_tensor * cpu_residual = + include_residual ? ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, output_size, token_count) : nullptr; + ggml_tensor * cpu_norm_weight = + include_next_rmsnorm ? ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, output_size) : nullptr; + ggml_tensor * cpu_projection = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); + ggml_tensor * cpu_postops = cpu_projection; + if (include_bias) { + cpu_postops = ggml_add(cpu_ctx, cpu_postops, cpu_bias); + } + if (include_residual) { + cpu_postops = ggml_add(cpu_ctx, cpu_postops, cpu_residual); + } + ggml_tensor * cpu_output = cpu_postops; + if (include_next_rmsnorm) { + ggml_tensor * cpu_rms = ggml_rms_norm(cpu_ctx, cpu_postops, kQwenRmsNormEps); + ggml_tensor * cpu_normalized = ggml_mul(cpu_ctx, cpu_rms, cpu_norm_weight); + cpu_output = ggml_add(cpu_ctx, cpu_normalized, cpu_postops); + } + + ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, input_size, output_size); + ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * hrx_bias = include_bias ? ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, output_size) : nullptr; + ggml_tensor * hrx_residual = + include_residual ? ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, output_size, token_count) : nullptr; + ggml_tensor * hrx_norm_weight = + include_next_rmsnorm ? ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, output_size) : nullptr; + ggml_tensor * hrx_projection = ggml_mul_mat(hrx_ctx, hrx_weight, hrx_input); + ggml_tensor * hrx_postops = hrx_projection; + if (include_bias) { + hrx_postops = ggml_add(hrx_ctx, hrx_postops, hrx_bias); + } + if (include_residual) { + hrx_postops = ggml_add(hrx_ctx, hrx_postops, hrx_residual); + } + ggml_tensor * hrx_output = hrx_postops; + if (include_next_rmsnorm) { + ggml_tensor * hrx_rms = ggml_rms_norm(hrx_ctx, hrx_postops, kQwenRmsNormEps); + ggml_tensor * hrx_normalized = ggml_mul(hrx_ctx, hrx_rms, hrx_norm_weight); + hrx_output = ggml_add(hrx_ctx, hrx_normalized, hrx_postops); + } REQUIRE(cpu_output != nullptr); REQUIRE(hrx_output != nullptr); @@ -1604,23 +4641,47 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); + if (include_next_rmsnorm) { + require_kernel_subsequence( + scheduled_kernel_sequence(hrx_graph), + { expected_kernel, "loom_libs:ggml_rmsnorm_binary_f32", "loom_libs:ggml_binary_f32" }); + } else { + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); + } ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - const std::vector weight = make_quantized_rows(weight_type, input_size, output_size, 6); - const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); + const std::vector weight(ggml_nbytes(cpu_weight), 0); + const std::vector input(static_cast(input_size * token_count), 0.00390625f); + const std::vector bias = make_pattern_f32(static_cast(output_size), 13, 0.02f); + const std::vector residual = make_pattern_f32(static_cast(output_size * token_count), 23, 0.02f); + const std::vector norm_weight = make_weight(output_size); set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + if (include_bias) { + set_tensor_pair_bytes(cpu_backend, cpu_bias, hrx_backend, hrx_bias, bias.data(), bias.size() * sizeof(float)); + } + if (include_residual) { + set_tensor_pair_bytes(cpu_backend, cpu_residual, hrx_backend, hrx_residual, residual.data(), + residual.size() * sizeof(float)); + } + if (include_next_rmsnorm) { + set_tensor_pair_bytes(cpu_backend, cpu_norm_weight, hrx_backend, hrx_norm_weight, norm_weight.data(), + norm_weight.size() * sizeof(float)); + } REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 3.0e-1f, 2.0e-2f); + if (!include_next_rmsnorm) { + require_close(get_f32_tensor(hrx_backend, hrx_postops), get_f32_tensor(cpu_backend, cpu_postops), 1.0f, + 3.0e-2f); + } + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0f, 5.0e-2f); ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); @@ -1630,6 +4691,48 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, ggml_backend_free(hrx_backend); } +static void run_dense_matmul_zero_weight_case(ggml_type weight_type, + const char * expected_kernel, + int64_t token_count = 2) { + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(32 * 1024 * 1024); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + constexpr int64_t input_size = kQwenHiddenSize; + constexpr int64_t output_size = 128; + 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); + 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); + require_kernel_subsequence(scheduled_kernel_sequence(graph), { expected_kernel }); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, hrx_backend); + REQUIRE(buffer != nullptr); + + const std::vector weight_bytes(ggml_nbytes(weight), 0); + const std::vector input_data = make_pattern_f32(input_size * token_count, 7, 0.01f); + const std::vector expected(static_cast(output_size * token_count), 0.0f); + set_tensor_bytes(hrx_backend, weight, weight_bytes.data(), weight_bytes.size()); + set_tensor_bytes(hrx_backend, input, input_data.data(), input_data.size() * sizeof(float)); + + REQUIRE(ggml_backend_graph_compute(hrx_backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(hrx_backend); + require_close(get_f32_tensor(hrx_backend, output), expected, 0.0f, 0.0f); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(hrx_backend); +} + static void run_endpoint_rmsnorm_q6k_q8_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -1696,6 +4799,681 @@ static void run_endpoint_rmsnorm_q6k_q8_cpu_reference_case() { ggml_backend_free(hrx_backend); } +static void run_lfm_head_rope_cpu_reference_case(int64_t head_count, int64_t token_count, bool cache_output) { + constexpr int64_t head_size = 64; + constexpr float epsilon = 1.0e-5f; + const int64_t row_size = head_size * head_count; + const int64_t cache_rows = token_count + 64; + + 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 = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + struct HeadGraph { + ggml_tensor * input; + ggml_tensor * weight; + ggml_tensor * positions; + ggml_tensor * cache; + ggml_tensor * row_ids; + ggml_tensor * output; + }; + + auto build = [&](ggml_context * ctx) { + HeadGraph g = {}; + g.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); + g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size); + g.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); + ggml_tensor * normalized = ggml_rms_norm(ctx, g.input, epsilon); + ggml_tensor * scaled = ggml_mul(ctx, normalized, g.weight); + ggml_tensor * rope = ggml_rope_ext(ctx, scaled, g.positions, nullptr, head_size, GGML_ROPE_TYPE_NEOX, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + g.output = rope; + if (cache_output) { + g.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, row_size, cache_rows); + g.row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * rows = ggml_reshape_2d(ctx, rope, row_size, token_count); + ggml_tensor * view = ggml_view_2d(ctx, rows, row_size, token_count, rows->nb[1], 0); + g.output = ggml_set_rows(ctx, g.cache, view, g.row_ids); + } + REQUIRE(g.output != nullptr); + return g; + }; + + HeadGraph cpu = build(cpu_ctx); + HeadGraph hrx = build(hrx_ctx); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu.output); + ggml_build_forward_expand(hrx_graph, hrx.output); + const std::vector kernels = scheduled_kernel_sequence(hrx_graph); + if (cache_output) { + REQUIRE(std::find(kernels.begin(), kernels.end(), "loom_libs:ggml_rmsnorm_mul_rope_f32") == kernels.end()); + require_kernel_subsequence(kernels, { "loom_libs:ggml_rope_f32", "loom_libs:ggml_set_rows" }); + } else { + require_kernel_subsequence(kernels, { "loom_libs:ggml_rmsnorm_mul_rope_f32" }); + } + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(row_size * token_count), 47, 0.03f); + const std::vector weight = make_weight(head_size); + std::vector positions(static_cast(token_count)); + std::vector row_ids(static_cast(token_count)); + for (int64_t i = 0; i < token_count; ++i) { + positions[static_cast(i)] = static_cast(3 + 2 * i); + row_ids[static_cast(i)] = (17 * i + 5) % cache_rows; + } + set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.weight, hrx_backend, hrx.weight, weight.data(), + weight.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.positions, hrx_backend, hrx.positions, positions.data(), + positions.size() * sizeof(int32_t)); + if (cache_output) { + const std::vector cache(static_cast(row_size * cache_rows), ggml_fp32_to_fp16(-3.0f)); + set_tensor_pair_bytes(cpu_backend, cpu.cache, hrx_backend, hrx.cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu.row_ids, hrx_backend, hrx.row_ids, row_ids.data(), + row_ids.size() * sizeof(int64_t)); + } + + 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); + const float tolerance = cache_output ? 1.0e-3f : 1.0e-4f; + require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), tolerance, + tolerance); + + 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_lfm_head_rope_availability_checks() { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 64, 8, 1); + ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 64); + ggml_tensor * pos_f32 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + ggml_tensor * positions = ggml_cast(ctx, pos_f32, GGML_TYPE_I32); + ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-5f); + ggml_tensor * scaled = ggml_mul(ctx, rms, weight); + ggml_tensor * output = ggml_rope_ext(ctx, scaled, positions, nullptr, 64, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, + 0.0f, 1.0f, 0.0f, 0.0f); + 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()); + const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); + REQUIRE(registry != nullptr); + const size_t rms_index = producer_index_for_tensor(imported.graph, rms); + const size_t pos_index = producer_index_for_tensor(imported.graph, positions); + ggml::hrx::CommandPlan plan; + std::vector covered(imported.graph.nodes().size(), false); + auto matched_kernel = [&]() { + ggml::hrx::DispatchMatch match; + const ggml::hrx::DispatchMatchContext context = { + imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, + plan, next_plan_value(imported.graph, plan), + }; + REQUIRE(registry->match(context, match)); + REQUIRE(match.dispatches.size() == 1); + return kernel_name_for_id(match.dispatches[0].kernel.kernel_id); + }; + REQUIRE(matched_kernel() == "loom_libs:ggml_rmsnorm_f32"); + covered[pos_index] = true; + REQUIRE(matched_kernel() == "loom_libs:ggml_rmsnorm_mul_rope_f32"); + + ggml_free(ctx); +} + +static void run_rope_set_rows_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); + + constexpr int64_t head_size = 8; + constexpr int64_t head_count = 2; + constexpr int64_t token_count = 3; + constexpr int64_t cache_row_count = 8; + constexpr int64_t hidden_size = head_size * head_count; + + { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * cpu_pos = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_freq = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * cpu_out = ggml_rope_ext(cpu_ctx, cpu_input, cpu_pos, cpu_freq, head_size, GGML_ROPE_TYPE_NEOX, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * hrx_input = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * hrx_pos = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_freq = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * hrx_out = ggml_rope_ext(hrx_ctx, hrx_input, hrx_pos, hrx_freq, head_size, GGML_ROPE_TYPE_NEOX, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_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); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(hidden_size * token_count), 41, 0.03f); + const std::vector pos = { 0, 1, 7 }; + const std::vector freq = { 1.0f, 0.5f, 0.25f, 0.125f }; + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), + input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_pos, hrx_backend, hrx_pos, pos.data(), pos.size() * sizeof(int32_t)); + set_tensor_pair_bytes(cpu_backend, cpu_freq, hrx_backend, hrx_freq, freq.data(), freq.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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-4f, 1.0e-4f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 8 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t partial_head_size = 128; + constexpr int64_t partial_n_dims = 96; + constexpr int64_t partial_head_count = 3; + constexpr int64_t partial_tokens = 4; + ggml_tensor * cpu_input = + ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, partial_head_size, partial_head_count, partial_tokens); + ggml_tensor * cpu_pos = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, partial_tokens); + ggml_tensor * cpu_freq = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, partial_n_dims / 2); + ggml_tensor * cpu_out = ggml_rope_ext(cpu_ctx, cpu_input, cpu_pos, cpu_freq, partial_n_dims, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * hrx_input = + ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, partial_head_size, partial_head_count, partial_tokens); + ggml_tensor * hrx_pos = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, partial_tokens); + ggml_tensor * hrx_freq = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, partial_n_dims / 2); + ggml_tensor * hrx_out = ggml_rope_ext(hrx_ctx, hrx_input, hrx_pos, hrx_freq, partial_n_dims, + GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rope_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = + make_pattern_f32(static_cast(partial_head_size * partial_head_count * partial_tokens), 47, 0.01f); + const std::vector pos = { 0, 1, 13, 29 }; + std::vector freq(static_cast(partial_n_dims / 2)); + for (size_t i = 0; i < freq.size(); ++i) { + freq[i] = 1.0f + 0.01f * static_cast(i); + } + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), + input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_pos, hrx_backend, hrx_pos, pos.data(), pos.size() * sizeof(int32_t)); + set_tensor_pair_bytes(cpu_backend, cpu_freq, hrx_backend, hrx_freq, freq.data(), freq.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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-4f, 1.0e-4f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + } + + auto run_view_set_rows_case = [&](int64_t view_token_count) { + ggml_init_params params = {}; + params.mem_size = 16 * 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + constexpr int64_t phi_hidden_size = 1024; + constexpr int64_t phi_cache_row_count = 256; + constexpr int64_t phi_row_stride = 5120; + constexpr size_t row_offset = 8 * sizeof(float); + const size_t storage_elements = + row_offset / sizeof(float) + static_cast((view_token_count - 1) * phi_row_stride + phi_hidden_size); + + ggml_tensor * cpu_cache = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F16, phi_hidden_size, phi_cache_row_count); + ggml_tensor * cpu_store = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * cpu_rows = ggml_view_2d(cpu_ctx, cpu_store, phi_hidden_size, view_token_count, + phi_row_stride * sizeof(float), row_offset); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I64, view_token_count); + ggml_tensor * cpu_out = ggml_set_rows(cpu_ctx, cpu_cache, cpu_rows, cpu_ids); + ggml_tensor * hrx_cache = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F16, phi_hidden_size, phi_cache_row_count); + ggml_tensor * hrx_store = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, storage_elements); + ggml_tensor * hrx_rows = ggml_view_2d(hrx_ctx, hrx_store, phi_hidden_size, view_token_count, + phi_row_stride * sizeof(float), row_offset); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I64, view_token_count); + ggml_tensor * hrx_out = ggml_set_rows(hrx_ctx, hrx_cache, hrx_rows, hrx_ids); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_set_rows" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector cache(static_cast(phi_hidden_size * phi_cache_row_count), + ggml_fp32_to_fp16(-2.5f)); + std::vector storage(storage_elements, -17.0f); + const std::vector rows = + make_pattern_f32(static_cast(phi_hidden_size * view_token_count), 52, 0.004f); + for (int64_t token = 0; token < view_token_count; ++token) { + std::memcpy(storage.data() + row_offset / sizeof(float) + static_cast(token * phi_row_stride), + rows.data() + static_cast(token * phi_hidden_size), + static_cast(phi_hidden_size) * sizeof(float)); + } + std::vector ids(static_cast(view_token_count)); + for (int64_t i = 0; i < view_token_count; ++i) { + ids[static_cast(i)] = (i * 17 + 3) % phi_cache_row_count; + } + + set_tensor_pair_bytes(cpu_backend, cpu_cache, hrx_backend, hrx_cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu_store, hrx_backend, hrx_store, storage.data(), + storage.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int64_t)); + + 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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-3f, 1.0e-3f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + }; + + run_view_set_rows_case(2); + run_view_set_rows_case(14); + + { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * cpu_pos = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_freq = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * cpu_out = ggml_rope_ext(cpu_ctx, cpu_input, cpu_pos, cpu_freq, head_size, GGML_ROPE_TYPE_NORMAL, + 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * hrx_input = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * hrx_pos = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_freq = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * hrx_out = ggml_rope_ext(hrx_ctx, hrx_input, hrx_pos, hrx_freq, head_size, GGML_ROPE_TYPE_NORMAL, + 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rope_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(hidden_size * token_count), 45, 0.03f); + const std::vector pos = { 0, 1, 7 }; + const std::vector freq = { 1.0f, 0.5f, 0.25f, 0.125f }; + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), + input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_pos, hrx_backend, hrx_pos, pos.data(), pos.size() * sizeof(int32_t)); + set_tensor_pair_bytes(cpu_backend, cpu_freq, hrx_backend, hrx_freq, freq.data(), freq.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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-4f, 1.0e-4f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_cache = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * cpu_rows = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F16, hidden_size, token_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * cpu_out = ggml_set_rows(cpu_ctx, cpu_cache, cpu_rows, cpu_ids); + ggml_tensor * hrx_cache = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * hrx_rows = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F16, hidden_size, token_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * hrx_out = ggml_set_rows(hrx_ctx, hrx_cache, hrx_rows, hrx_ids); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_set_rows" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector cache(static_cast(hidden_size * cache_row_count), + ggml_fp32_to_fp16(-4.0f)); + const std::vector rows = + make_pattern_f16(static_cast(hidden_size * token_count), 44, 0.04f); + const std::vector ids = { 7, 0, 5 }; + set_tensor_pair_bytes(cpu_backend, cpu_cache, hrx_backend, hrx_cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu_rows, hrx_backend, hrx_rows, rows.data(), + rows.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int64_t)); + + 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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-3f, 1.0e-3f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_cache = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * cpu_rows = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * cpu_out = ggml_set_rows(cpu_ctx, cpu_cache, cpu_rows, cpu_ids); + ggml_tensor * hrx_cache = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * hrx_rows = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * hrx_out = ggml_set_rows(hrx_ctx, hrx_cache, hrx_rows, hrx_ids); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_set_rows" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector cache(static_cast(hidden_size * cache_row_count), + ggml_fp32_to_fp16(-2.0f)); + const std::vector rows = make_pattern_f32(static_cast(hidden_size * token_count), 42, 0.04f); + const std::vector ids = { 6, 2, 4 }; + set_tensor_pair_bytes(cpu_backend, cpu_cache, hrx_backend, hrx_cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu_rows, hrx_backend, hrx_rows, rows.data(), rows.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int64_t)); + + 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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-3f, 1.0e-3f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * cpu_pos = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_freq = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * cpu_rope = ggml_rope_ext(cpu_ctx, cpu_input, cpu_pos, cpu_freq, head_size, GGML_ROPE_TYPE_NEOX, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * cpu_rows = ggml_reshape_2d(cpu_ctx, cpu_rope, hidden_size, token_count); + ggml_tensor * cpu_cache = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * cpu_out = ggml_set_rows(cpu_ctx, cpu_cache, cpu_rows, cpu_ids); + ggml_tensor * hrx_input = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * hrx_pos = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_freq = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * hrx_rope = ggml_rope_ext(hrx_ctx, hrx_input, hrx_pos, hrx_freq, head_size, GGML_ROPE_TYPE_NEOX, 0, + 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * hrx_rows = ggml_reshape_2d(hrx_ctx, hrx_rope, hidden_size, token_count); + ggml_tensor * hrx_cache = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * hrx_out = ggml_set_rows(hrx_ctx, hrx_cache, hrx_rows, hrx_ids); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rope_set_rows_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(hidden_size * token_count), 43, 0.03f); + const std::vector pos = { 0, 1, 7 }; + const std::vector freq = { 1.0f, 0.5f, 0.25f, 0.125f }; + const std::vector cache(static_cast(hidden_size * cache_row_count), + ggml_fp32_to_fp16(-3.0f)); + const std::vector ids = { 5, 1, 3 }; + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), + input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_pos, hrx_backend, hrx_pos, pos.data(), pos.size() * sizeof(int32_t)); + set_tensor_pair_bytes(cpu_backend, cpu_freq, hrx_backend, hrx_freq, freq.data(), freq.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_cache, hrx_backend, hrx_cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int64_t)); + + 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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-3f, 1.0e-3f); + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + } + + { + ggml_init_params params = {}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + ggml_tensor * cpu_input = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * cpu_pos = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * cpu_freq = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * cpu_rope = ggml_rope_ext(cpu_ctx, cpu_input, cpu_pos, cpu_freq, head_size, GGML_ROPE_TYPE_NORMAL, + 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * cpu_rows = ggml_reshape_2d(cpu_ctx, cpu_rope, hidden_size, token_count); + ggml_tensor * cpu_cache = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * cpu_out = ggml_set_rows(cpu_ctx, cpu_cache, cpu_rows, cpu_ids); + ggml_tensor * hrx_input = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, head_size, head_count, token_count); + ggml_tensor * hrx_pos = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); + ggml_tensor * hrx_freq = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, head_size / 2); + ggml_tensor * hrx_rope = ggml_rope_ext(hrx_ctx, hrx_input, hrx_pos, hrx_freq, head_size, GGML_ROPE_TYPE_NORMAL, + 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + ggml_tensor * hrx_rows = ggml_reshape_2d(hrx_ctx, hrx_rope, hidden_size, token_count); + ggml_tensor * hrx_cache = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F16, hidden_size, cache_row_count); + ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I64, token_count); + ggml_tensor * hrx_out = ggml_set_rows(hrx_ctx, hrx_cache, hrx_rows, hrx_ids); + REQUIRE(cpu_out != nullptr); + REQUIRE(hrx_out != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_out); + ggml_build_forward_expand(hrx_graph, hrx_out); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rope_set_rows_f32" }); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + const std::vector input = make_pattern_f32(static_cast(hidden_size * token_count), 46, 0.03f); + const std::vector pos = { 0, 1, 7 }; + const std::vector freq = { 1.0f, 0.5f, 0.25f, 0.125f }; + const std::vector cache(static_cast(hidden_size * cache_row_count), + ggml_fp32_to_fp16(-5.0f)); + const std::vector ids = { 4, 2, 6 }; + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), + input.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_pos, hrx_backend, hrx_pos, pos.data(), pos.size() * sizeof(int32_t)); + set_tensor_pair_bytes(cpu_backend, cpu_freq, hrx_backend, hrx_freq, freq.data(), freq.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_cache, hrx_backend, hrx_cache, cache.data(), + cache.size() * sizeof(ggml_fp16_t)); + set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int64_t)); + + 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_out), get_f32_tensor(cpu_backend, cpu_out), 1.0e-3f, 1.0e-3f); + + 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_attention_postprocess_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -1720,6 +5498,10 @@ static void run_attention_postprocess_cpu_reference_case() { AttentionPostprocessGraph hrx = build_attention_postprocess_graph(hrx_ctx, token_count, query_head_count, key_value_head_count, cache_row_count); + auto hrx_query_buffer = allocate_test_weight(hrx_backend, hrx.query_weight); + auto hrx_key_buffer = allocate_test_weight(hrx_backend, hrx.key_weight); + auto hrx_value_buffer = allocate_test_weight(hrx_backend, hrx.value_weight); + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); REQUIRE(cpu_graph != nullptr); @@ -1821,11 +5603,10 @@ static void run_routed_moe_cpu_reference_case(ggml_type down_weight_type, bool i std::vector expected = { "qwen3_moe:qwen3_moe_router_top8_f32", - "qwen3_moe:qwen3_moe_build_expert_table", - "qwen3_moe:qwen3_moe_build_expert_partition_table", - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_f16_wmma", - down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_down_q4k_f16_wmma_grouped" : - "qwen3_moe:qwen3_moe_routed_down_q6k_f16_wmma_grouped", + "loom_libs:ggml_moe_build_expert_table", + "loom_libs:ggml_moe_build_expert_partition_table", + "loom_libs:ggml_mul_mat_id_swiglu_f16_f16_wmma", + "loom_libs:ggml_mul_mat_id_f16_f16_wmma", include_next_rmsnorm ? "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32" : "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32", }; @@ -1953,18 +5734,22 @@ 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", "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", - down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8" : - "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8", - down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_down_q4k_q8_1_x4_next_q8" : - "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8", + down_weight_type == GGML_TYPE_Q4_K ? "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_q8" : + "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_f32", + down_weight_type == GGML_TYPE_Q4_K ? "loom_libs:ggml_mul_mat_id_vector_weighted_wave32_publish_q8" : + "loom_libs:ggml_mul_mat_id_vector_weighted_wave64_publish_q8", }; require_kernel_subsequence(scheduled_kernel_sequence(graph), expected); ggml_free(ctx); @@ -1990,39 +5775,789 @@ static void run_decode_attention_qkv_scheduling_case() { ggml_free(ctx); } -int main() { - run_rmsnorm_support_checks(); - run_alternate_value_alias_lookup_checks(); +struct SsmConvPrefillGraph { + ggml_context * ctx = nullptr; + ggml_tensor * input = nullptr; + ggml_tensor * weight = nullptr; + ggml_tensor * pool = nullptr; + ggml_tensor * initial = nullptr; + ggml_tensor * ids = nullptr; + ggml_tensor * filter = nullptr; + ggml_tensor * output = nullptr; + ggml_tensor * cache = nullptr; + ggml_cgraph * graph = nullptr; +}; + +static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, + int64_t k, + int64_t n, + bool expose, + bool gathered, + bool clear_state) { + ggml_init_params params = {}; + params.mem_size = 32 * 1024 * 1024; + params.no_alloc = true; + SsmConvPrefillGraph c; + c.ctx = ggml_init(params); + REQUIRE(c.ctx != nullptr); + c.input = ggml_new_tensor_2d(c.ctx, GGML_TYPE_F32, k, 512); + c.weight = ggml_new_tensor_2d(c.ctx, type, k, n); + c.pool = ggml_new_tensor_2d(c.ctx, GGML_TYPE_F32, 3 * n, 2); + c.initial = ggml_new_tensor_3d(c.ctx, GGML_TYPE_F32, 3, n, 1); + c.ids = ggml_new_tensor_1d(c.ctx, GGML_TYPE_I32, 1); + c.filter = ggml_new_tensor_2d(c.ctx, GGML_TYPE_F32, 4, n); + c.graph = ggml_new_graph(c.ctx); + auto state = c.initial; + if (gathered) { + if (clear_state) { + auto row = ggml_view_2d(c.ctx, c.pool, 3 * n, 1, c.pool->nb[1], 0); + ggml_build_forward_expand(c.graph, ggml_scale_inplace(c.ctx, row, 0.0f)); + } + state = ggml_get_rows(c.ctx, c.pool, c.ids); + ggml_build_forward_expand(c.graph, state); + state = ggml_reshape_3d(c.ctx, state, 3, n, 1); + } + auto raw = ggml_mul_mat(c.ctx, c.weight, c.input); + auto x = ggml_reshape_3d(c.ctx, raw, n, 512, 1); + auto window = ggml_concat(c.ctx, state, ggml_transpose(c.ctx, x), 0); + auto tail = ggml_view_3d(c.ctx, window, 3, n, 1, window->nb[1], window->nb[2], 512 * sizeof(float)); + c.cache = ggml_view_2d(c.ctx, c.pool, 3 * n, 1, c.pool->nb[1], 0); + ggml_build_forward_expand(c.graph, ggml_cpy(c.ctx, tail, c.cache)); + c.output = ggml_silu(c.ctx, ggml_ssm_conv(c.ctx, window, c.filter)); + ggml_set_output(c.output); + ggml_build_forward_expand(c.graph, c.output); + if (expose) { + ggml_set_output(raw); + ggml_build_forward_expand(c.graph, ggml_scale(c.ctx, raw, 0.5f)); + } + const auto sequence = scheduled_kernel_sequence(c.graph); + const char * fused = gathered ? "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" : + "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4"; + REQUIRE((std::find(sequence.begin(), sequence.end(), fused) != sequence.end()) == !expose); + if (gathered && !expose) { + REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32") != + sequence.end()); + } + return c; +} + +static void run_ssm_conv_prefill_recurrent_case(ggml_type type, bool gathered, bool clear_state) { + const int64_t k = 256; + const int64_t n = 8192; + auto backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + auto a = make_ssm_conv_prefill_graph(type, k, n, false, gathered, clear_state); + auto b = make_ssm_conv_prefill_graph(type, k, n, true, gathered, clear_state); + auto aw = allocate_test_weight(backend, a.weight); + auto bw = allocate_test_weight(backend, b.weight); + auto ab = ggml_backend_alloc_ctx_tensors(a.ctx, backend); + auto bb = ggml_backend_alloc_ctx_tensors(b.ctx, backend); + REQUIRE(ab != nullptr && bb != nullptr); + const auto weights = make_matmul_weight_bytes(type, k, n, 6); + const auto filters = make_pattern_f32(4 * n, 11, 0.0625f); + const auto states = make_pattern_f32(6 * n, 17, 0.03125f); + set_tensor_pair_bytes(backend, a.weight, backend, b.weight, weights.data(), weights.size()); + set_tensor_pair_bytes(backend, a.filter, backend, b.filter, filters.data(), filters.size() * sizeof(float)); + set_tensor_pair_bytes(backend, a.pool, backend, b.pool, states.data(), states.size() * sizeof(float)); + set_tensor_pair_bytes(backend, a.initial, backend, b.initial, states.data(), 3 * n * sizeof(float)); + for (int round = 0; round < 5; ++round) { + const int32_t id = round % 2; + set_tensor_pair_bytes(backend, a.ids, backend, b.ids, &id, sizeof(id)); + const auto input = make_pattern_f32(k * 512, 7 + round, 0.0625f); + set_tensor_pair_bytes(backend, a.input, backend, b.input, input.data(), input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, a.graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(backend, b.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const auto actual = get_f32_tensor(backend, a.output); + const auto expected = get_f32_tensor(backend, b.output); + require_close(get_f32_tensor(backend, a.cache), get_f32_tensor(backend, b.cache), 0.0f, 0.0f); + require_close(actual, expected, 0.0f, 0.0f); + } + ggml_backend_buffer_free(ab); + ggml_backend_buffer_free(bb); + aw.reset(); + bw.reset(); + ggml_free(a.ctx); + ggml_free(b.ctx); + ggml_backend_free(backend); +} + +static void run_vector_output_capacity_experiment() { + constexpr int64_t kSmallVectorInputSize = 256; + constexpr int64_t kAboveOldOutputSize = 262208; + constexpr int64_t kLiftedOutputSize = 524288; + + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, + kSmallVectorInputSize); + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, + kAboveOldOutputSize, kSmallVectorInputSize); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kAboveOldOutputSize, true, false, false, kSmallVectorInputSize); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kLiftedOutputSize, + kSmallVectorInputSize); +} + +static void run_vector_postops_cpu_reference_case() { + constexpr int64_t kVectorInputSize = 256; + constexpr int64_t kVectorOutputSize = 256; + + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kVectorOutputSize, true, false, false, kVectorInputSize); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kVectorOutputSize, true, true, false, kVectorInputSize); +} + +static const char * glu_op_name(ggml_glu_op op) { + switch (op) { + case GGML_GLU_OP_REGLU: + return "reglu"; + case GGML_GLU_OP_GEGLU: + return "geglu"; + case GGML_GLU_OP_SWIGLU: + return "swiglu"; + case GGML_GLU_OP_SWIGLU_OAI: + return "swiglu_oai"; + case GGML_GLU_OP_GEGLU_ERF: + return "geglu_erf"; + case GGML_GLU_OP_GEGLU_QUICK: + return "geglu_quick"; + case GGML_GLU_OP_COUNT: + break; + } + return "unknown_glu"; +} + +static std::string type_name(ggml_type type) { + return ggml_type_name(type); +} + +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); }); + suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".gathered", + [type] { run_ssm_conv_prefill_recurrent_case(type, true, false); }); + suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".gathered_clear_state", + [type] { run_ssm_conv_prefill_recurrent_case(type, true, true); }); + } + suite.device_case("view_mul_f32.cpu_reference", [] { run_view_mul_f32_cpu_reference_case(); }); + for (const int64_t token_count : { 2, 23 }) { + suite.device_case("strided_cont.tokens" + std::to_string(token_count), + [token_count] { run_strided_cont_cpu_reference_case(token_count); }); + } + for (const ggml_glu_op op : + { GGML_GLU_OP_REGLU, GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { + suite.device_case(std::string("glu_split_f32.") + glu_op_name(op), + [op] { run_glu_split_f32_cpu_reference_case(op); }); + suite.device_case(std::string("glu_packed_f32.") + glu_op_name(op) + ".normal", + [op] { run_glu_packed_f32_cpu_reference_case(op, false); }); + } + suite.device_case("glu_packed_f32.swiglu.gated", + [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, true); }); + suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens2", + [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 2); }); + suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens14", + [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 14); }); + suite.device_case("gather_add_f32.cpu_reference", [] { run_gather_add_f32_cpu_reference_case(); }); + suite.device_case("scale_add_f32.cpu_reference.packed", [] { run_scale_add_f32_cpu_reference_case(false); }); + suite.device_case("scale_add_f32.cpu_reference.strided", [] { run_scale_add_f32_cpu_reference_case(true); }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_XS, GGML_TYPE_F16, + GGML_TYPE_BF16, GGML_TYPE_F32 }) { + suite.device_case("get_rows_f32." + type_name(type), [type] { run_get_rows_f32_cpu_reference_case(type); }); + } + suite.device_case("get_rows_f32.q1_0.rows2048", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q1_0, 2048); }); + suite.device_case("get_rows_f32.q5_1.rows640", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); + suite.device_case("get_rows_q8_1.zero_weight", [] { run_get_rows_q8_1_zero_weight_case(); }); + suite.device_case("get_rows_scale_f32.cpu_reference", [] { run_get_rows_scale_f32_cpu_reference_case(); }); + suite.device_case("get_rows_scale_f32.q5_1.cpu_reference", + [] { run_get_rows_scale_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); +} + +static void register_dense_matmul_cases(Suite & suite) { + // K=384 exercises a partial 256-weight tile; N=65 exercises an unpaired output row. + for (const int64_t input_size : { 256, 384, 768 }) { + suite.device_case("dense_matmul.q1_0.vector_tail.k" + std::to_string(input_size), [input_size] { + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 65, input_size); + }); + } + + 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( + "dense_matmul.q4_k.tokens" + std::to_string(tokens) + ".outputs" + std::to_string(outputs), + [tokens, outputs] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, + tokens == 1 ? + "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + tokens, outputs); + }); + } + } + for (const ggml_type type : + { GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { + suite.device_case("dense_matmul." + type_name(type) + ".tokens3.outputs47", [type] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 47); + }); + } + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + suite.device_case("dense_matmul." + type_name(type) + ".tiled.tokens33.outputs65.input512", [type] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 33, 65, + 512); + }); + } + suite.device_case("dense_matmul_unary.f32.tokens3.outputs47", [] { + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + 3, 47); + }); + suite.device_case("get_rows_f32.f32.rows262144.tokens4.outputs1", + [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); }); + suite.device_case("dense_matmul.q4_k.skinny.tokens2.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, + 128); + }); + suite.device_case("dense_matmul.q4_k.vector.tokens1.outputs128", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); + }); + 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, + kQwenHiddenSize, true); + }); + // Exercise the production K dimension and a partial output tile. + suite.device_case("dense_matmul.iq4_xs.prefill.tokens512.outputs64", [] { + run_dense_matmul_cpu_reference_case( + GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 512, 64, 3072, true); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + for (const ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { + suite.device_case( + "attention_v_matmul." + type_name(type) + ".cache_" + type_name(cache_type), [type, cache_type] { + run_dense_matmul_cpu_reference_case(type, + "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32", 1, + 1024, 2048, false, false, cache_type); + }); + } + } + suite.device_case("attention_v_matmul.q6_k.prefill_f16.repeated_indices", + [] { run_prefill_value_cache_repeated_indices_case(); }); + for (const ggml_type type : + { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_IQ2_S, + GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, GGML_TYPE_F16, GGML_TYPE_BF16 }) { + const char * vector_kernel = type == GGML_TYPE_Q6_K ? "loom_libs:ggml_mul_mat_vector_q6_f32_f32" : + "loom_libs:ggml_mul_mat_vector_f32_f32"; + const int64_t outputs = type == GGML_TYPE_IQ2_S ? 640 : (type == GGML_TYPE_IQ4_NL ? 1024 : 128); + const int64_t input = type == GGML_TYPE_IQ4_NL ? 640 : kQwenHiddenSize; + const char * route = type == GGML_TYPE_IQ4_NL ? "tiled" : "skinny"; + const char * kernel = type == GGML_TYPE_IQ4_NL ? "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; + suite.device_case( + "dense_matmul." + type_name(type) + "." + route + ".tokens2.outputs" + std::to_string(outputs), + [type, outputs, input, kernel] { run_dense_matmul_cpu_reference_case(type, kernel, 2, outputs, input); }); + suite.device_case("dense_matmul." + type_name(type) + ".vector.tokens1.outputs" + std::to_string(outputs), + [type, vector_kernel, outputs, input] { + run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, outputs, input); + }); + } + suite.device_case("dense_matmul.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); + }); + suite.device_case("dense_matmul.q5_0.tiled.tokens2.outputs128.input544", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, + 128, 544); + }); + suite.device_case("dense_matmul.q5_0.vector.tokens1.outputs256.input640", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 640); + }); + suite.device_case("dense_matmul.q8_1.zero_weight.skinny", [] { + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"); + }); + suite.device_case("dense_matmul.q8_1.zero_weight.vector", [] { + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1); + }); + suite.device_case("dense_matmul.f32.skinny.tokens2.outputs256", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, + 256); + }); + suite.device_case("dense_matmul.f32.vector.tokens1.outputs256", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256); + }); + suite.device_case("dense_matmul_unary.f32.tokens2.outputs128", [] { + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + 2, 128); + }); + suite.device_case("dense_matmul_swiglu.q4_k.f16.swiglu.tokens2.outputs128", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + }); + suite.device_case("dense_matmul_swiglu.q4_k.f16.geglu.tokens2.outputs128", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, + GGML_GLU_OP_GEGLU); + }); + suite.device_case("dense_matmul_swiglu.iq4_nl.iq4_nl.geglu.tokens2.outputs2048.input640", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, + 2048, GGML_GLU_OP_GEGLU, 640); + }); + 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( + GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, op); + }); + } + suite.device_case("dense_matmul_packed_glu.q4_k.tokens2.outputs128", [] { + run_dense_matmul_packed_glu_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128); + }); + suite.device_case("dense_matmul_packed_glu.q4_k.swiglu.tokens2.outputs128.gated", [] { + run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + 2, 128, GGML_GLU_OP_SWIGLU, true); + }); + suite.device_case("dense_matmul_binary.q4_k.f16.tokens2.outputs128", [] { + run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + }); + suite.device_case("vector_q8_publish.cpu_reference", [] { run_vector_q8_publish_cpu_reference_case(); }); + for (const ggml_glu_op op : + { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { + for (const int64_t outputs : { 128, 4096 }) { + suite.device_case(std::string("dense_matmul_swiglu.q4_k.q4_k.lowtoken.") + glu_op_name(op) + + ".tokens4.outputs" + std::to_string(outputs), + [op, outputs] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, outputs, op); + }); + } + for (const int64_t tokens : { 256, 512 }) { + if (tokens == 512) { + suite.device_case( + std::string("dense_matmul_swiglu.q4_k.q4_k.prefill.") + glu_op_name(op) + ".tokens512.outputs4096", + [op] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 512, 4096, op); + }); + } else { + for (const int64_t outputs : { 128, 4096 }) { + suite.device_case(std::string("dense_matmul_swiglu.q4_k.q4_k.prefill.") + glu_op_name(op) + + ".tokens256.outputs" + std::to_string(outputs), + [op, outputs] { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, + outputs, op); + }); + } + } + } + } + for (const int64_t tokens : { 1, 2, 3, 5 }) { + suite.device_case( + "dense_matmul_swiglu.q4_k.q4_k.lowtoken.tokens" + std::to_string(tokens) + ".outputs4160", [tokens] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 4160); + }); + } + for (const int64_t tokens : { 1, 2, 5 }) { + suite.device_case( + "dense_matmul_swiglu.q4_k.q4_k.swiglu.tokens" + std::to_string(tokens) + ".outputs16384.input4096", + [tokens] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); + }); + } + suite.device_case("dense_matmul_swiglu.q4_k.q4_k.reglu.tokens3.outputs16384.input4352", [] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 3, 16384, + GGML_GLU_OP_REGLU, 4352); + }); + suite.device_case("dense_matmul_binary.q4_k.q4_k.lowtoken.tokens4.outputs128", [] { + run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128); + }); + suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs128", [] { + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 128); + }); + suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs4096", [] { + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 4096); + }); + suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens1024.outputs2048", [] { + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, 2048); + }); + suite.device_case("dense_matmul_postops.f16.bias.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, true, false, false); + }); + suite.device_case("dense_matmul_postops.f16.residual.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, false, true, false); + }); + suite.device_case("dense_matmul_postops.f16.bias_residual.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, true, true, false); + }); + suite.device_case("dense_matmul_postops.f16.residual.alias.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, false, true, true); + }); + suite.device_case("dense_matmul_postops.f16.bias_residual.alias.tokens33.outputs256", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", + 33, 256, true, true, true); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + const int64_t input_size = type == GGML_TYPE_Q4_K ? 20480 : 8192; + suite.device_case( + "dense_matmul." + type_name(type) + ".skinny.tokens5.outputs4096.input" + std::to_string(input_size), + [type, input_size] { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, + 4096, input_size); + }); + suite.device_case("dense_matmul_postops." + type_name(type) + ".add.tokens5.outputs4160", [type, input_size] { + run_dense_matmul_postops_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", 5, 4160, false, true, false, + input_size + 256); + }); + suite.device_case("dense_matmul_postops." + type_name(type) + ".add.alias.tokens5.outputs4224", + [type, input_size] { + run_dense_matmul_postops_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", 5, 4224, + false, true, true, input_size + 256); + }); + } + suite.device_case("dense_matmul.q4_k.skinny.tokens5.outputs4096.input8192", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, + 4096, 8192); + }); + for (const int64_t input_size : { 6144, 6400, 6656 }) { + suite.device_case( + "dense_matmul.q4_k.skinny.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 5120, input_size); + }); + suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens5.outputs5120.input" + std::to_string(input_size), + [input_size] { + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 5, 5120, false, true, true, input_size); + }); + } + suite.device_case("dense_matmul.q4_k.skinny.tokens4.outputs4096.input4096", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, + 4096, 4096); + }); + suite.device_case("dense_matmul_postops.q4_k.add.tokens4.outputs4160.input6400", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 4, 4160, false, true, false, 6400); + }); + suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens4.outputs4224.input6656", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 4, 4224, false, true, true, 6656); + }); + suite.device_case("endpoint_rmsnorm_q6k_q8.cpu_reference", + [] { run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); }); + suite.device_case("dense_matmul.q6_k.skinny.tokens5.outputs4096.input16384", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, + 4096, 16384); + }); + suite.device_case("dense_matmul_postops.q6_k.add.tokens5.outputs4160.input16640", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 5, 4160, false, true, false, 16640); + }); + suite.device_case("dense_matmul_postops.q6_k.add.alias.tokens5.outputs4224.input16640", [] { + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, + "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", + 5, 4224, false, true, true, 16640); + }); + for (const int64_t tokens : { 1, 3, 5 }) { + suite.device_case( + "dense_matmul.q4_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + tokens, 256, 5120, true); + }); + suite.device_case( + "dense_matmul.q6_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + tokens, 256, 5120, true, true); + }); + } + suite.device_case("dense_matmul.q6_k.vector_q6.tokens1.outputs256.input5120", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 256, 5120, + true); + }); + suite.device_case("dense_matmul.q6_k.vector_q6.tokens1.outputs320.input512", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 320, 512, + true); + }); +} + +static void register_rmsnorm_and_scheduling_cases(Suite & suite) { + suite.host_case("support.lfm_head_rope.availability", [] { run_lfm_head_rope_availability_checks(); }); + suite.host_case("gemma_post_proj.matcher_negatives", [] { run_gemma_post_proj_matcher_negatives(); }); + for (const int64_t token_count : { 1, 3, 64 }) { + suite.device_case("gemma_post_proj.tokens" + std::to_string(token_count), + [token_count] { run_gemma_post_proj_cpu_reference_case(token_count); }); + } + suite.device_case("rmsnorm.default", [] { run_rmsnorm_cpu_reference_case(); }); + suite.device_case("rmsnorm.hidden3840.tokens18", [] { run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); }); + for (const int64_t token_count : { 2, 23 }) { + suite.device_case("strided_rmsnorm.tokens" + std::to_string(token_count), + [token_count] { run_strided_rmsnorm_cpu_reference_case(token_count); }); + } + suite.device_case("strided_rmsnorm_binary.minicpm_pp512_witness", + [] { run_strided_rmsnorm_binary_cpu_reference_case(); }); + suite.device_case("rmsnorm_binary_add.cpu_reference", [] { run_rmsnorm_binary_add_cpu_reference_case(); }); + suite.device_case("rmsnorm_gate.silu.hidden128.tokens1.stride7", + [] { run_rmsnorm_gate_cpu_reference_case(128, 1, 7, GGML_UNARY_OP_SILU); }); + suite.device_case("rmsnorm_gate.gelu.hidden256.tokens3.stride3", + [] { run_rmsnorm_gate_cpu_reference_case(256, 3, 3, GGML_UNARY_OP_GELU); }); + suite.device_case("rmsnorm_gate.relu.hidden512.tokens3.stride5", + [] { run_rmsnorm_gate_cpu_reference_case(512, 3, 5, GGML_UNARY_OP_RELU); }); + suite.device_case("rmsnorm_gate.gelu_erf.hidden1024.tokens1.stride4", + [] { run_rmsnorm_gate_cpu_reference_case(1024, 1, 4, GGML_UNARY_OP_GELU_ERF); }); + suite.device_case("rmsnorm_gate.silu.hidden128.tokens48.stride512", + [] { run_rmsnorm_gate_cpu_reference_case(128, 48, 512, GGML_UNARY_OP_SILU); }); + suite.device_case("rmsnorm_mul.hidden3840.tokens18", [] { run_rmsnorm_mul_cpu_reference_case(3840, 18, 1.0e-6f); }); + for (const int64_t hidden_size : { 1024, 2048, 2560, 3072, 4096, 5120 }) { + suite.device_case("rmsnorm_mul.lowtoken.random.hidden" + std::to_string(hidden_size) + ".tokens1", + [hidden_size] { run_rmsnorm_mul_cpu_reference_case(hidden_size, 1, 1.0e-6f); }); + suite.device_case("rmsnorm_mul.lowtoken.adversarial.hidden" + std::to_string(hidden_size) + ".tokens1", + [hidden_size] { run_rmsnorm_mul_cpu_reference_case(hidden_size, 1, 1.0e-5f, true); }); + } + for (const int64_t hidden_size : { 2048, 5120 }) { + for (const int64_t token_count : { 2, 4 }) { + suite.device_case("rmsnorm_mul.lowtoken.random.hidden" + std::to_string(hidden_size) + ".tokens" + + std::to_string(token_count), + [hidden_size, token_count] { + run_rmsnorm_mul_cpu_reference_case(hidden_size, token_count, 1.0e-6f); + }); + } + } + suite.device_case("gemma_scaled_rmsnorm.cpu_reference", [] { run_gemma_scaled_rmsnorm_cpu_reference_case(); }); + suite.device_case("gemma_scaled_rmsnorm_mul.cpu_reference", + [] { run_gemma_scaled_rmsnorm_mul_cpu_reference_case(); }); + suite.device_case("scheduled.scale_rmsnorm", [] { run_scheduled_hrx_scale_rmsnorm_case(); }); + suite.device_case("scheduled.rmsnorm_scale", [] { run_scheduled_hrx_rmsnorm_scale_case(); }); + suite.device_case("scheduled.rmsnorm_view_scale", [] { run_scheduled_hrx_rmsnorm_view_scale_case(); }); + suite.device_case("external_view_input_rmsnorm", [] { run_external_view_input_rmsnorm_case(); }); + suite.device_case("split_local_view_alias_import", [] { run_split_local_view_alias_import_case(); }); + suite.device_case("rope_set_rows.cpu_reference", [] { run_rope_set_rows_cpu_reference_case(); }); + for (const int64_t head_count : { 8, 32 }) { + for (const int64_t token_count : { 1, 64 }) { + suite.device_case( + "lfm_head_rope.query.heads" + std::to_string(head_count) + ".tokens" + std::to_string(token_count), + [head_count, token_count] { run_lfm_head_rope_cpu_reference_case(head_count, token_count, false); }); + } + } + suite.device_case("lfm_head_rope.query.heads7.tokens3", [] { run_lfm_head_rope_cpu_reference_case(7, 3, false); }); + for (const int64_t token_count : { 1, 64 }) { + suite.device_case("lfm_head_rope.cache.heads8.tokens" + std::to_string(token_count), + [token_count] { run_lfm_head_rope_cpu_reference_case(8, token_count, true); }); + } + suite.device_case("attention_postprocess.cpu_reference", [] { run_attention_postprocess_cpu_reference_case(); }); + suite.device_case("routed_moe.q4_k.next_rmsnorm", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); }); + suite.device_case("flash_attention.asymmetric", [] { run_asymmetric_flash_attention_cpu_reference_case(); }); + suite.device_case("routed_moe.q6_k", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); }); + suite.device_case("decode_attention_qkv.scheduling", [] { run_decode_attention_qkv_scheduling_case(); }); + suite.device_case("decode_routed_moe.q4_k", [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); }); + suite.device_case("decode_routed_moe.q6_k", [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); }); + suite.device_case("decode_routed_moe.q6_k.alias_gate_input", + [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); }); + suite.device_case("rmsnorm_mul.hidden256.tokens1", [] { run_rmsnorm_mul_case(256, 1); }); + suite.device_case("rmsnorm_mul.hidden256.tokens4", [] { run_rmsnorm_mul_case(256, 4); }); + suite.device_case("rmsnorm_mul.hidden2048.tokens1", [] { run_rmsnorm_mul_case(2048, 1); }); + suite.device_case("router_projection.tokens4", [] { run_router_projection_case(4); }); + suite.device_case("router_top8.tokens4", [] { run_router_top8_case(4); }); + suite.device_case("qwen_flash_attention", [] { run_qwen_flash_attention_case(); }); +} + +static void register_flash_attention_cases(Suite & suite) { + 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] { + run_common_flash_attention_cpu_reference_case( + head_size, 2, 65, 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + }); + suite.device_case( + "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query16.kv64", [head_size] { + run_common_flash_attention_cpu_reference_case(head_size, 16, 64, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_f32_f16_wmma"); + }); + suite.device_case( + "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query17.kv65", [head_size] { + run_common_flash_attention_cpu_reference_case(head_size, 17, 65, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_f32_f16_wmma"); + }); + } + suite.device_case("common_flash_attention.decode.head96.query2.kv4.capacity320", [] { + run_common_flash_attention_cpu_reference_case( + 96, 2, 4, 1.0f / std::sqrt(96.0f), "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 320); + }); + suite.device_case("common_flash_attention.decode.head320.query2.kv64.capacity64", [] { + run_common_flash_attention_cpu_reference_case( + 320, 2, 64, 1.0f / std::sqrt(320.0f), "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", + 64); + }); + suite.device_case("qwen_decode_split_flash_attention.query1.kv512", + [] { run_qwen_decode_split_flash_attention_scheduling_case(1, 512); }); + suite.device_case("qwen_decode_split_flash_attention.query4.kv513", + [] { run_qwen_decode_split_flash_attention_scheduling_case(4, 513); }); + suite.device_case("qwen_decode_split_flash_attention.query1.kv512.masked", + [] { run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true); }); + suite.device_case("qwen_decode_split_flash_attention.query1.kv512.masked.head96.value64", [] { + run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true, 96, 64, 1.0f / std::sqrt(96.0f)); + }); + suite.device_case("qwen_decode_split_flash_attention.query4.kv513.masked.head96.value64", [] { + run_qwen_decode_split_flash_attention_scheduling_case(4, 513, true, 96, 64, 1.0f / std::sqrt(96.0f)); + }); + suite.device_case("qwen_decode_attention_output_next_q8.no_selectors", + [] { run_qwen_decode_attention_output_next_q8_scheduling_case(false); }); + suite.device_case("qwen_decode_attention_output_next_q8.selectors", + [] { run_qwen_decode_attention_output_next_q8_scheduling_case(true); }); + suite.device_case("qwen_full_cache_prefill_flash_attention.tokens512", + [] { run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); }); +} + +static void register_hrx_ops_cases(Suite & suite) { + register_basic_ops_cases(suite); + register_dense_matmul_cases(suite); + register_rmsnorm_and_scheduling_cases(suite); + register_flash_attention_cases(suite); +} + +int main(int argc, char ** argv) { + Suite suite(test_runner::Config::with_prefix("HRX test", "hrx", "GGML_HRX_TEST")); + register_hrx_ops_cases(suite); - if (ggml_backend_hrx_get_device_count() == 0) { - std::fprintf(stderr, "test skipped: no HRX devices available\n"); + test_runner::Options options; + if (!suite.parse_options(argc, argv, options)) { + suite.print_usage(argc > 0 ? argv[0] : "test-hrx-ops"); + return 2; + } + if (options.help) { + return suite.run(options, false); + } + + const bool run_vector_cap_experiment = std::getenv("GGML_HRX_RUN_VECTOR_CAP_EXPERIMENT") != nullptr; + const bool run_vector_publish_test = std::getenv("GGML_HRX_RUN_VECTOR_PUBLISH_TEST") != nullptr; + const bool run_vector_postops_test = std::getenv("GGML_HRX_RUN_VECTOR_POSTOPS_TEST") != nullptr; + + const bool has_device = ggml_backend_hrx_get_device_count() != 0; + if (!options.runner_control && run_vector_cap_experiment) { + run_vector_output_capacity_experiment(); + return 0; + } + if (!options.runner_control && run_vector_publish_test) { + run_vector_q8_publish_cpu_reference_case(); + return 0; + } + if (!options.runner_control && run_vector_postops_test) { + run_vector_postops_cpu_reference_case(); return 0; } - run_add_f32_cpu_reference_case(); - run_gather_add_f32_cpu_reference_case(); - run_token_embedding_q4k_cpu_reference_case(); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "qwen3_moe:qwen3_moe_dense_linear_q4k_f16_wmma", 2, 128); - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "qwen3_moe:qwen3_moe_dense_linear_q6k_f16_wmma", 2, 128); - run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); - run_attention_postprocess_cpu_reference_case(); - run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); - run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); - run_decode_attention_qkv_scheduling_case(); - run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); - run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); - run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); - run_rmsnorm_mul_case(256, 1); - run_rmsnorm_mul_case(256, 4); - run_rmsnorm_mul_case(2048, 1); - run_router_projection_case(4); - run_router_top8_case(4); - run_qwen_flash_attention_case(); - run_qwen_decode_split_flash_attention_scheduling_case(1, 512); - run_qwen_decode_split_flash_attention_scheduling_case(4, 513); - run_qwen_decode_attention_output_next_q8_scheduling_case(false); - run_qwen_decode_attention_output_next_q8_scheduling_case(true); - run_qwen_full_cache_prefill_flash_attention_scheduling_case(16); - run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); - return 0; + return suite.run(options, has_device); } diff --git a/tests/testing_suite.h b/tests/testing_suite.h new file mode 100644 index 000000000000..ffa5fa0e66cb --- /dev/null +++ b/tests/testing_suite.h @@ -0,0 +1,548 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#endif + +namespace test_runner { + +struct Case { + std::string name; + std::function run; + bool requires_device = true; +}; + +enum class OutputMode { + Failed, + All, + None, +}; + +struct Config { + std::string suite_name = "test"; + std::string list_flag = "--test-list"; + std::string case_flag = "--test-case"; + std::string filter_flag = "--test-filter"; + std::string jobs_flag = "--test-jobs"; + std::string output_flag = "--test-output"; + std::string env_jobs; + std::string env_filter; + std::string env_output; + size_t max_default_jobs = 4; + + static Config with_prefix(std::string suite_name, + const std::string & flag_prefix, + const std::string & env_prefix, + size_t max_default_jobs = 4) { + Config config; + config.suite_name = std::move(suite_name); + config.list_flag = "--" + flag_prefix + "-test-list"; + config.case_flag = "--" + flag_prefix + "-test-case"; + config.filter_flag = "--" + flag_prefix + "-test-filter"; + config.jobs_flag = "--" + flag_prefix + "-test-jobs"; + config.output_flag = "--" + flag_prefix + "-test-output"; + config.env_jobs = env_prefix + "_JOBS"; + config.env_filter = env_prefix + "_FILTER"; + config.env_output = env_prefix + "_OUTPUT"; + config.max_default_jobs = max_default_jobs; + return config; + } +}; + +struct Options { + std::string executable; + std::string single_case; + std::string filter; + size_t jobs = 0; + bool help = false; + bool list = false; + bool runner_control = false; + OutputMode output = OutputMode::Failed; +}; + +struct Result { + std::string name; + std::string output; + int exit_code = -1; + int signal_code = 0; + int64_t elapsed_ms = 0; + bool passed = false; +}; + +struct Selection { + std::vector cases; + size_t matched = 0; + size_t skipped_device = 0; +}; + +inline void add_case(std::vector & cases, + std::string name, + std::function run, + bool requires_device = true) { + cases.push_back({ std::move(name), std::move(run), requires_device }); +} + +inline size_t default_jobs(const Config & config) { + const unsigned int hardware_jobs = std::thread::hardware_concurrency(); + if (hardware_jobs == 0) { + return 1; + } + return std::min(config.max_default_jobs, hardware_jobs); +} + +inline bool parse_positive_size(const char * value, size_t & out) { + if (value == nullptr || *value == '\0') { + return false; + } + char * end = nullptr; + errno = 0; + const unsigned long parsed = std::strtoul(value, &end, 10); + if (errno != 0 || end == value || *end != '\0' || parsed == 0) { + return false; + } + out = static_cast(parsed); + return true; +} + +inline bool parse_output_mode(const char * value, OutputMode & out) { + if (value == nullptr) { + return true; + } + const std::string mode = value; + if (mode == "failed" || mode == "fail") { + out = OutputMode::Failed; + return true; + } + if (mode == "all") { + out = OutputMode::All; + return true; + } + if (mode == "none") { + out = OutputMode::None; + return true; + } + return false; +} + +inline void print_usage(const Config & config, const char * executable) { + std::fprintf(stderr, + "usage: %s [%s N] [%s REGEX] [%s MODE] [%s]\n" + " %s %s NAME\n" + "\n" + "environment: %s, %s, %s\n" + "output modes: failed, all, none\n", + executable, config.jobs_flag.c_str(), config.filter_flag.c_str(), config.output_flag.c_str(), + config.list_flag.c_str(), executable, config.case_flag.c_str(), config.env_jobs.c_str(), + config.env_filter.c_str(), config.env_output.c_str()); +} + +inline bool parse_options(int argc, char ** argv, const Config & config, Options & options) { + options.executable = argc > 0 ? argv[0] : config.suite_name.c_str(); + if (!config.env_filter.empty()) { + if (const char * env_filter = std::getenv(config.env_filter.c_str())) { + options.filter = env_filter; + options.runner_control = true; + } + } + if (!config.env_jobs.empty()) { + if (const char * env_jobs = std::getenv(config.env_jobs.c_str())) { + if (!parse_positive_size(env_jobs, options.jobs)) { + std::fprintf(stderr, "invalid %s value: %s\n", config.env_jobs.c_str(), env_jobs); + return false; + } + options.runner_control = true; + } + } + if (!config.env_output.empty()) { + if (const char * env_output = std::getenv(config.env_output.c_str())) { + if (!parse_output_mode(env_output, options.output)) { + std::fprintf(stderr, "invalid %s value: %s\n", config.env_output.c_str(), env_output); + return false; + } + options.runner_control = true; + } + } + + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == config.list_flag) { + options.list = true; + options.runner_control = true; + } else if (arg == config.case_flag) { + if (++i >= argc) { + std::fprintf(stderr, "%s requires a case name\n", config.case_flag.c_str()); + return false; + } + options.single_case = argv[i]; + options.runner_control = true; + } else if (arg == config.filter_flag) { + if (++i >= argc) { + std::fprintf(stderr, "%s requires a regex\n", config.filter_flag.c_str()); + return false; + } + options.filter = argv[i]; + options.runner_control = true; + } else if (arg == config.jobs_flag) { + if (++i >= argc || !parse_positive_size(argv[i], options.jobs)) { + std::fprintf(stderr, "%s requires a positive integer\n", config.jobs_flag.c_str()); + return false; + } + options.runner_control = true; + } else if (arg == config.output_flag) { + if (++i >= argc || !parse_output_mode(argv[i], options.output)) { + std::fprintf(stderr, "%s requires failed, all, or none\n", config.output_flag.c_str()); + return false; + } + options.runner_control = true; + } else if (arg == "--help" || arg == "-h") { + options.help = true; + options.runner_control = true; + } else { + std::fprintf(stderr, "unknown argument: %s\n", arg.c_str()); + return false; + } + } + + if (options.jobs == 0) { + options.jobs = default_jobs(config); + } + return true; +} + +inline const Case * find_case(const std::vector & cases, const std::string & name) { + for (const Case & test_case : cases) { + if (test_case.name == name) { + return &test_case; + } + } + return nullptr; +} + +inline int run_single_case_in_process(const Config & config, + const std::vector & cases, + const std::string & name, + bool has_device) { + const Case * test_case = find_case(cases, name); + if (test_case == nullptr) { + std::fprintf(stderr, "unknown test case: %s\n", name.c_str()); + return 2; + } + if (test_case->requires_device && !has_device) { + std::fprintf(stderr, "test skipped: no device available for %s\n", config.suite_name.c_str()); + return 0; + } + test_case->run(); + return 0; +} + +inline Result run_single_case_child(const Config & config, const Options & options, const Case & test_case) { + Result result; + result.name = test_case.name; + + const auto start = std::chrono::steady_clock::now(); + +#if defined(_WIN32) + result.output = "process-isolated test runner is not supported on this platform\n"; + result.exit_code = 127; + result.elapsed_ms = + std::chrono::duration_cast(std::chrono::steady_clock::now() - start).count(); + return result; +#else + int pipe_fds[2] = { -1, -1 }; + if (pipe(pipe_fds) != 0) { + result.output = "failed to create output pipe\n"; + result.exit_code = 127; + return result; + } + + std::fflush(nullptr); + const pid_t pid = fork(); + if (pid == 0) { + close(pipe_fds[0]); + dup2(pipe_fds[1], STDOUT_FILENO); + dup2(pipe_fds[1], STDERR_FILENO); + close(pipe_fds[1]); + + std::vector child_argv; + child_argv.push_back(const_cast(options.executable.c_str())); + child_argv.push_back(const_cast(config.case_flag.c_str())); + child_argv.push_back(const_cast(test_case.name.c_str())); + child_argv.push_back(nullptr); + execvp(child_argv[0], child_argv.data()); + std::fprintf(stderr, "failed to exec %s: %s\n", child_argv[0], std::strerror(errno)); + _exit(127); + } + + close(pipe_fds[1]); + if (pid < 0) { + close(pipe_fds[0]); + result.output = "failed to fork child test process\n"; + result.exit_code = 127; + return result; + } + + char buffer[4096]; + for (;;) { + const ssize_t bytes = read(pipe_fds[0], buffer, sizeof(buffer)); + if (bytes > 0) { + result.output.append(buffer, static_cast(bytes)); + continue; + } + if (bytes == 0) { + break; + } + if (errno == EINTR) { + continue; + } + result.output += "failed to read child output\n"; + break; + } + close(pipe_fds[0]); + + int status = 0; + while (waitpid(pid, &status, 0) < 0) { + if (errno == EINTR) { + continue; + } + result.output += "failed to wait for child test process\n"; + result.exit_code = 127; + return result; + } + + if (WIFEXITED(status)) { + result.exit_code = WEXITSTATUS(status); + result.passed = result.exit_code == 0; + } else if (WIFSIGNALED(status)) { + result.signal_code = WTERMSIG(status); + result.passed = false; + } + + result.elapsed_ms = + std::chrono::duration_cast(std::chrono::steady_clock::now() - start).count(); + return result; +#endif +} + +inline bool should_print_output(OutputMode mode, const Result & result) { + if (mode == OutputMode::All) { + return true; + } + if (mode == OutputMode::None) { + return false; + } + return !result.passed; +} + +inline int run_cases_parallel(const Config & config, + const Options & options, + const std::vector & selected_cases) { + if (selected_cases.empty()) { + std::printf("No %s cases selected\n", config.suite_name.c_str()); + return 0; + } + + const size_t jobs = std::min(options.jobs, selected_cases.size()); + std::vector results(selected_cases.size()); + std::atomic next_index(0); + std::atomic finished(0); + std::mutex print_mutex; + const auto suite_start = std::chrono::steady_clock::now(); + + std::printf("Running %zu %s case(s) with %zu worker process(es)\n", selected_cases.size(), + config.suite_name.c_str(), jobs); + + auto worker = [&] { + for (;;) { + const size_t index = next_index.fetch_add(1); + if (index >= selected_cases.size()) { + return; + } + Result result = run_single_case_child(config, options, selected_cases[index]); + const size_t done = finished.fetch_add(1) + 1; + { + std::lock_guard lock(print_mutex); + std::printf("%s %4zu/%zu %-86s %8lld ms\n", result.passed ? "[PASS]" : "[FAIL]", done, + selected_cases.size(), result.name.c_str(), static_cast(result.elapsed_ms)); + if (!result.passed) { + if (result.signal_code != 0) { + std::printf(" terminated by signal %d\n", result.signal_code); + } else { + std::printf(" exit code %d\n", result.exit_code); + } + } + if (should_print_output(options.output, result) && !result.output.empty()) { + std::printf("----- output: %s -----\n%s", result.name.c_str(), result.output.c_str()); + if (result.output.back() != '\n') { + std::printf("\n"); + } + std::printf("----- end output: %s -----\n", result.name.c_str()); + } + std::fflush(stdout); + } + results[index] = std::move(result); + } + }; + + std::vector workers; + workers.reserve(jobs); + for (size_t i = 0; i < jobs; ++i) { + workers.emplace_back(worker); + } + for (std::thread & thread : workers) { + thread.join(); + } + + size_t passed = 0; + for (const Result & result : results) { + if (result.passed) { + ++passed; + } + } + const size_t failed = results.size() - passed; + const int64_t elapsed_ms = + std::chrono::duration_cast(std::chrono::steady_clock::now() - suite_start).count(); + + std::printf("\n%s summary: %zu passed, %zu failed, %zu total, %lld ms\n", config.suite_name.c_str(), passed, failed, + results.size(), static_cast(elapsed_ms)); + if (failed != 0) { + std::printf("Failed cases:\n"); + for (const Result & result : results) { + if (!result.passed) { + std::printf(" %s\n", result.name.c_str()); + } + } + } + return failed == 0 ? 0 : 1; +} + +inline bool select_cases(const std::vector & cases, + const std::string & filter_text, + bool has_device, + bool include_unavailable_device_cases, + Selection & selection) { + try { + std::regex filter; + const bool has_filter = !filter_text.empty(); + if (has_filter) { + filter = std::regex(filter_text); + } + for (const Case & test_case : cases) { + if (has_filter && !std::regex_search(test_case.name, filter)) { + continue; + } + ++selection.matched; + if (!has_device && test_case.requires_device && !include_unavailable_device_cases) { + ++selection.skipped_device; + continue; + } + selection.cases.push_back(test_case); + } + } catch (const std::regex_error & error) { + std::fprintf(stderr, "invalid test filter regex: %s\n", error.what()); + return false; + } + return true; +} + +inline void print_case_list(const std::vector & cases) { + for (const Case & test_case : cases) { + std::printf("[%s] %s\n", test_case.requires_device ? "device" : "host", test_case.name.c_str()); + } +} + +inline bool validate_case_names(const std::vector & cases) { + for (size_t i = 0; i < cases.size(); ++i) { + for (size_t j = i + 1; j < cases.size(); ++j) { + if (cases[i].name == cases[j].name) { + std::fprintf(stderr, "duplicate test case: %s\n", cases[i].name.c_str()); + return false; + } + } + } + return true; +} + +class Suite { + public: + explicit Suite(Config config) : config_(std::move(config)) {} + + template void device_case(std::string name, F && run) { + add_case(cases_, std::move(name), std::function(std::forward(run)), true); + } + + template void host_case(std::string name, F && run) { + add_case(cases_, std::move(name), std::function(std::forward(run)), false); + } + + bool parse_options(int argc, char ** argv, Options & options) const { + return test_runner::parse_options(argc, argv, config_, options); + } + + void print_usage(const char * executable) const { test_runner::print_usage(config_, executable); } + + int run(int argc, char ** argv, bool has_device) const { + Options options; + if (!parse_options(argc, argv, options)) { + print_usage(argc > 0 ? argv[0] : config_.suite_name.c_str()); + return 2; + } + return run(options, has_device); + } + + int run(const Options & options, bool has_device) const { + if (options.help) { + print_usage(options.executable.c_str()); + return 0; + } + if (!validate_case_names(cases_)) { + return 2; + } + if (!options.single_case.empty()) { + return run_single_case_in_process(config_, cases_, options.single_case, has_device); + } + + Selection selection; + if (!select_cases(cases_, options.filter, has_device, options.list, selection)) { + return 2; + } + + if (options.list) { + print_case_list(selection.cases); + return 0; + } + + if (selection.skipped_device != 0) { + std::fprintf(stderr, "test skipped: no device available for %s (%zu device case(s))\n", + config_.suite_name.c_str(), selection.skipped_device); + if (selection.cases.empty()) { + return 0; + } + } + + return run_cases_parallel(config_, options, selection.cases); + } + + private: + Config config_; + std::vector cases_; +}; + +} // namespace test_runner